diff --git a/.circleci/config.yml b/.circleci/config.yml index 0d442e32b0a..1edf22578ff 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1291,7 +1291,7 @@ jobs: name: Install MCP SDK1 peer command: | uv venv --python 3.12 .venv-mcp-peer - uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' + uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - when: condition: diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index d2ab3478397..a33695d0dbe 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -33,7 +33,7 @@ while IFS= read -r file || [ -n "$file" ]; do *) has_backend=true ;; esac case "$file" in - model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json) + model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json) has_cost_map=true ;; tests/unit/proxy/*) : ;; *) outside_cost_map_set=true ;; diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a707eb8044c..f6c06f7d83f 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,4 +1,4 @@ -/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri -/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri +/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri +/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri /helm/ @yassin-berriai /terraform/ @yassin-berriai diff --git a/.github/scripts/auto_update_price_and_context_window_file.py b/.github/scripts/auto_update_price_and_context_window_file.py deleted file mode 100644 index 4c9a05bd4f0..00000000000 --- a/.github/scripts/auto_update_price_and_context_window_file.py +++ /dev/null @@ -1,319 +0,0 @@ -import asyncio -import aiohttp -import json -import math -from typing import Any - -# Asynchronously fetch data from a given URL -async def fetch_data(url): - try: - # Create an asynchronous session - async with aiohttp.ClientSession() as session: - # Send a GET request to the URL - async with session.get(url) as resp: - # Raise an error if the response status is not OK - resp.raise_for_status() - # Parse the response JSON - resp_json = await resp.json() - print("Fetch the data from URL.") - # Return the 'data' field from the JSON response - return resp_json['data'] - except Exception as e: - # Print an error message if fetching data fails - print("Error fetching data from URL:", e) - return None - - -FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models" -FRIENDLI_PROVIDER = "friendliai" - -INHERITABLE_BASE_KEYS = ( - "supports_pdf_input", - "supports_assistant_prefill", - "supports_adaptive_thinking", - "supports_output_config", -) - -REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max") - - -def _find_base_model_entry(base_model: str, local_data: dict) -> str | None: - if not base_model: - return None - bm_tail = base_model.split("/")[-1].lower() - if base_model in local_data: - return base_model - for key in local_data: - if key.startswith("sample_spec") or key == "fallback_generalizations": - continue - if key.split("/")[-1].lower() == bm_tail: - return key - return None - - -def _reasoning_effort_levels(reasoning_options: list) -> list: - offered = { - val - for opt in reasoning_options or [] - if opt.get("type") == "effort" - for val in opt.get("values", []) - } - return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered] - - -def _valid_token_price(value: object) -> bool: - try: - price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except - except (TypeError, ValueError): - return False - return math.isfinite(price) and price >= 0 - - -def _has_valid_token_prices(pricing: dict | None) -> bool: - prices = pricing or {} - return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output")) - - -def _pricing(pricing: dict) -> dict: - out: dict[str, Any] = {} - if not pricing: - return out - if "input" in pricing: - out["input_cost_per_token"] = float(pricing["input"]) - if "output" in pricing: - out["output_cost_per_token"] = float(pricing["output"]) - if "input_cache_read" in pricing and pricing["input_cache_read"] is not None: - out["cache_read_input_token_cost"] = float(pricing["input_cache_read"]) - return out - - -def _modality_flags(input_mods: list) -> dict: - mods = input_mods or [] - has_image = "image" in mods - return { - "supports_vision": has_image, - "supports_image_input": has_image, - "supports_video_input": "video" in mods, - } - - -def transform_friendli_data(data: list, local_data: dict) -> dict: - transformed: dict[str, dict] = {} - if not data: - return transformed - for model in data: - # An unpriced row must never wholesale-replace an already priced local entry: - # missing prices cost-calculate as zero, silently zeroing tracked spend - if not _has_valid_token_prices(model.get("pricing")): - continue - model_id = model["id"] - base_model = model.get("base_model") or "" - entry: dict[str, Any] = { - "litellm_provider": FRIENDLI_PROVIDER, - } - - base_key = _find_base_model_entry(base_model, local_data) - if base_key: - base_entry = local_data[base_key] - for k in INHERITABLE_BASE_KEYS: - if k in base_entry: - entry[k] = base_entry[k] - - ctx = model.get("context_length") - if ctx is not None: - entry["max_input_tokens"] = int(ctx) - max_out = model.get("max_completion_tokens") - if max_out is not None: - entry["max_output_tokens"] = int(max_out) - entry["max_tokens"] = int(max_out) - - pricing = _pricing(model.get("pricing", {})) - entry.update(pricing) - entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing - - reasoning = model.get("reasoning") is True - entry["supports_reasoning"] = reasoning - if reasoning: - entry["reasoning_effort_levels"] = _reasoning_effort_levels( - model.get("reasoning_options", []) - ) - - func = model.get("functionality", {}) - entry["supports_function_calling"] = func.get("tool_call") is True - entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True - is_struct = func.get("structured_output") is True - entry["supports_response_schema"] = is_struct - entry["supports_native_structured_output"] = is_struct - entry["supports_system_messages"] = func.get("system_messages") is True - entry["supports_tool_choice"] = func.get("tool_choice") is True - - entry.update(_modality_flags(model.get("input_modalities", []))) - - entry["mode"] = model.get("mode", "chat") - - desc = model.get("description") - if desc: - entry["comment"] = desc - - dep = model.get("deprecation_date") - if dep: - entry["deprecation_date"] = dep.split("T")[0] - - entry["source"] = FRIENDLI_API_URL - - transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry - return transformed - -# Synchronize local data with remote data -def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()): - # Update existing keys in local_data with values from remote_data - # (replace_keys entries are swapped wholesale so a field the remote catalog - # dropped, e.g. cache pricing, cannot survive as a stale value) - for key in (set(local_data) & set(remote_data)): - if key in replace_keys: - local_data[key] = remote_data[key] - else: - local_data[key].update(remote_data[key]) - - # Add new keys from remote_data to local_data - for key in (set(remote_data) - set(local_data)): - local_data[key] = remote_data[key] - -# Write data to the json file -def write_to_file(file_path, data): - try: - # Open the file in write mode - with open(file_path, "w") as file: - # Dump the data as JSON into the file - json.dump(data, file, indent=4) - print("Values updated successfully.") - except Exception as e: - # Print an error message if writing to file fails - print("Error updating JSON file:", e) - -# Update the existing models and add the missing models for OpenRouter -def transform_openrouter_data(data): - transformed = {} - if not data: - return transformed - for row in data: - # Add the fields 'max_tokens' and 'input_cost_per_token' - obj = { - "max_tokens": row["context_length"], - "input_cost_per_token": float(row["pricing"]["prompt"]), - } - - # Add 'max_output_tokens' as a field if it is not None - if "top_provider" in row and "max_completion_tokens" in row["top_provider"] and row["top_provider"]["max_completion_tokens"] is not None: - obj['max_output_tokens'] = int(row["top_provider"]["max_completion_tokens"]) - - # Add the field 'output_cost_per_token' - obj.update({ - "output_cost_per_token": float(row["pricing"]["completion"]), - }) - - # Add field 'input_cost_per_image' if it exists and is non-zero - if "pricing" in row and "image" in row["pricing"] and float(row["pricing"]["image"]) != 0.0: - obj['input_cost_per_image'] = float(row["pricing"]["image"]) - - # Add the fields 'litellm_provider' and 'mode' - obj.update({ - "litellm_provider": "openrouter", - "mode": "chat" - }) - - # Add the 'supports_vision' field if the modality is 'multimodal' - if row.get('architecture', {}).get('modality') == 'multimodal': - obj['supports_vision'] = True - - # Use a composite key to store the transformed object - transformed[f'openrouter/{row["id"]}'] = obj - - return transformed - -# Update the existing models and add the missing models for Vercel AI Gateway -def transform_vercel_ai_gateway_data(data): - transformed = {} - if not data: - return transformed - for row in data: - # Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync - if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any( - row.get("pricing", {}).get(k) is None for k in ("input", "output") - ): - continue - obj = { - "max_tokens": row["context_window"], - "input_cost_per_token": float(row["pricing"]["input"]), - "output_cost_per_token": float(row["pricing"]["output"]), - 'max_output_tokens': row['max_tokens'], - 'max_input_tokens': row["context_window"], - } - - # Handle cache pricing if available - if "pricing" in row: - if "input_cache_read" in row["pricing"] and row["pricing"]["input_cache_read"] is not None: - obj['cache_read_input_token_cost'] = float(f"{float(row['pricing']['input_cache_read']):e}") - - if "input_cache_write" in row["pricing"] and row["pricing"]["input_cache_write"] is not None: - obj['cache_creation_input_token_cost'] = float(f"{float(row['pricing']['input_cache_write']):e}") - - mode = "embedding" if "embedding" in row["id"].lower() else "chat" - - obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode}) - - transformed[f'vercel_ai_gateway/{row["id"]}'] = obj - - return transformed - - -# Load local data from a specified file -def load_local_data(file_path): - try: - # Open the file in read mode - with open(file_path, "r") as file: - # Load and return the JSON data - return json.load(file) - except FileNotFoundError: - # Print an error message if the file is not found - print("File not found:", file_path) - return None - except json.JSONDecodeError as e: - # Print an error message if JSON decoding fails - print("Error decoding JSON:", e) - return None - -def main(): - local_file_path = "model_prices_and_context_window.json" # Path to the local data file - openrouter_url = "https://openrouter.ai/api/v1/models" # URL to fetch OpenRouter data - vercel_ai_gateway_url = "https://ai-gateway.vercel.sh/v1/models" # URL to fetch Vercel AI Gateway data - - # Load local data from file - local_data = load_local_data(local_file_path) - - # Fetch OpenRouter data - openrouter_data = asyncio.run(fetch_data(openrouter_url)) - # Transform the fetched OpenRouter data - openrouter_data = transform_openrouter_data(openrouter_data) - - # Fetch Vercel AI Gateway data - vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url)) - # Transform the fetched Vercel AI Gateway data - vercel_data = transform_vercel_ai_gateway_data(vercel_data) - - friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL)) - friendli_data = transform_friendli_data(friendli_data, local_data) - - # Combine both datasets - all_remote_data = {**openrouter_data, **vercel_data, **friendli_data} - - # If both local and openrouter data are available, synchronize and save - if local_data and all_remote_data: - sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data)) - write_to_file(local_file_path, local_data) - else: - print("Failed to fetch model data from either local file or URL.") - -# Entry point of the script -if __name__ == "__main__": - main() diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml deleted file mode 100644 index 838ba35f831..00000000000 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ /dev/null @@ -1,36 +0,0 @@ -name: Updates model_prices_and_context_window.json and Create Pull Request - -on: - schedule: - - cron: "0 0 * * 0" # Run every Sundays at midnight - #- cron: "0 0 * * *" # Run daily at midnight - -permissions: - contents: write - pull-requests: write - -jobs: - auto_update_price_and_context_window: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - 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: Create Pull Request - run: | - 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" \ - --head auto-update-price-and-context-window-$(date +'%Y-%m-%d') \ - --base main - env: - GH_TOKEN: ${{ secrets.GH_TOKEN }} diff --git a/.github/workflows/interactions-spec-canary.yml b/.github/workflows/interactions-spec-canary.yml new file mode 100644 index 00000000000..7ec3defdcb7 --- /dev/null +++ b/.github/workflows/interactions-spec-canary.yml @@ -0,0 +1,91 @@ +name: Interactions spec canary + +on: + schedule: + - cron: "17 7 * * *" + pull_request: + paths: + - .github/workflows/interactions-spec-canary.yml + - tests/unit/interactions/** + - litellm/types/interactions/** + workflow_dispatch: + +permissions: {} + +concurrency: interactions-spec-canary + +jobs: + interactions-spec-canary: + name: interactions-spec-canary + if: github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 15 + permissions: + contents: read + outputs: + compliance: ${{ steps.compliance.outcome }} + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project + + - name: Run the compliance tests against Google's live spec + id: compliance + run: | + uv run --no-sync pytest tests/unit/interactions/test_openapi_compliance.py \ + --interactions-spec-url=https://ai.google.dev/static/api/interactions.openapi.json \ + --tb=short -vv -p no:rerunfailures -p no:pytest-retry \ + --allow-hosts=127.0.0.1,::1,localhost,ai.google.dev + + tracking-issue: + name: tracking-issue + needs: interactions-spec-canary + if: always() && github.event_name != 'pull_request' && github.ref_name == github.event.repository.default_branch && github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + issues: write + steps: + - name: File the tracking issue unless the compliance step passed, close it when it did + uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + env: + COMPLIANCE_OUTCOME: ${{ needs.interactions-spec-canary.outputs.compliance }} + RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }} + with: + script: | + const title = "Interactions spec canary is red"; + const outcome = process.env.COMPLIANCE_OUTCOME; + const runUrl = process.env.RUN_URL; + const query = `repo:${context.repo.owner}/${context.repo.repo} is:issue is:open in:title "${title}"`; + const { data: found } = await github.rest.search.issuesAndPullRequests({ q: query, per_page: 20 }); + const open = found.items.find((issue) => issue.title === title); + if (outcome === "success") { + if (open === undefined) return; + await github.rest.issues.createComment({ ...context.repo, issue_number: open.number, body: `Green again: ${runUrl}` }); + await github.rest.issues.update({ ...context.repo, issue_number: open.number, state: "closed" }); + return; + } + const state = outcome === "failure" ? "red" : `not run (compliance step outcome: ${outcome || "none"})`; + if (open !== undefined) { + await github.rest.issues.createComment({ ...context.repo, issue_number: open.number, body: `Still ${state}: ${runUrl}` }); + return; + } + await github.rest.issues.create({ + ...context.repo, + title, + assignees: ["mateo-berri"], + body: `The scheduled run of tests/unit/interactions/test_openapi_compliance.py against Google's live Interactions spec is ${state}: ${runUrl}\n\nlitellm/types/interactions/README.md says how to refresh the snapshot and regenerate the types.`, + }); diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml deleted file mode 100644 index 5325510e45e..00000000000 --- a/.github/workflows/sync-together-ai-models.yml +++ /dev/null @@ -1,65 +0,0 @@ -name: Sync Together AI model registry - -on: - schedule: - - cron: "30 6 * * *" - workflow_dispatch: - -permissions: - contents: write - pull-requests: write - -jobs: - sync_together_ai_models: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - env: - BASE_BRANCH: ${{ github.event.repository.default_branch }} - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - ref: ${{ env.BASE_BRANCH }} - persist-credentials: false - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - name: Look for an already-open sync PR - id: existing - run: | - open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \ - --jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')" - echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT" - if [ -n "$open_pr" ]; then - echo "An open sync PR already exists on branch $open_pr; skipping this run." - fi - env: - GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} - - name: Run the sync - if: steps.existing.outputs.open_pr == '' - run: | - 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: Create a pull request when the registry changed - if: steps.existing.outputs.open_pr == '' - run: | - if git diff --quiet; then - echo "Registry already in sync; no PR needed." - exit 0 - fi - branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')" - git config user.name "github-actions[bot]" - 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 - git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')" - gh auth setup-git - git push origin "$branch" - gh pr create --title "feat(models): sync together_ai model registry" \ - --body-file "$RUNNER_TEMP/pr_body.md" \ - --head "$branch" \ - --base "$BASE_BRANCH" - env: - GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} diff --git a/.github/workflows/test-mcp-dependency-resolution.yml b/.github/workflows/test-mcp-dependency-resolution.yml index 3f0d1260e49..20efb351985 100644 --- a/.github/workflows/test-mcp-dependency-resolution.yml +++ b/.github/workflows/test-mcp-dependency-resolution.yml @@ -45,7 +45,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"] steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 @@ -132,7 +132,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"] steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: diff --git a/.github/workflows/test-merge-smoke.yml b/.github/workflows/test-merge-smoke.yml index 0d1cc6596de..f3f09523491 100644 --- a/.github/workflows/test-merge-smoke.yml +++ b/.github/workflows/test-merge-smoke.yml @@ -33,7 +33,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"] env: LITELLM_LOCAL_MODEL_COST_MAP: "True" steps: @@ -76,6 +76,12 @@ jobs: - name: Import and CLI checks run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py cli + - name: Routing and UUID compatibility + run: >- + $IN_NETNS .venv/bin/python -m pytest + tests/unit/router_strategy/test_complexity_router.py::TestCustomDimensions + tests/unit/test_uuid_helper.py -q + - name: Proxy startup check run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py proxy-startup --diagnostics-dir "${RUNNER_TEMP}/smoke-diagnostics" diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index c593acbe8fb..76af64c8ccd 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -216,10 +216,12 @@ jobs: - shard: llms-providers test-path: |- + tests/unit/llms/test_api_base_trailing_slash_routes.py tests/unit/llms/test_cache_control_and_reasoning.py tests/unit/llms/test_custom_llm.py tests/unit/llms/test_file_content_block.py tests/unit/llms/test_file_search_responses.py + tests/unit/llms/test_get_complete_url_trailing_slash.py tests/unit/llms/test_lifecycle_fix.py tests/unit/llms/test_oss_decision.py tests/unit/llms/test_polling_url_origin_match.py @@ -423,6 +425,7 @@ jobs: tests/unit/realtime_api tests/unit/rerank_api tests/unit/search + tests/unit/text_completions tests/unit/vector_stores tests/unit/videos python-version: "3.12" diff --git a/Dockerfile b/Dockerfile index 890c90fe449..4cef63c8d71 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -61,6 +61,7 @@ RUN apk add --no-cache \ python-3.13-dev \ rust \ openssl \ + libssl3 \ openssl-dev \ nodejs \ npm \ @@ -120,8 +121,8 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \ FROM $LITELLM_BUILD_IMAGE AS liteadmin-builder COPY --from=uvbin /uv /usr/local/bin/uv RUN apk add --no-cache python-3.13 -ADD --checksum=sha256:2f7ae5cdd9d91731c0990e74a58239dc3e3fd2bf28dab23b55eafcdc47aaf87e \ - https://github.com/BerriAI/litellm-admin-agent/archive/ef501e94bc9fbacb9233b922abf71427f030408c.tar.gz /tmp/liteadmin.tar.gz +ADD --checksum=sha256:4fb1847e28e2be30ba76dac94f02c4903291326516a0f4ba772f6ffb8b25d05c \ + https://github.com/BerriAI/litellm-admin-agent/archive/5661c805ce479a0cf11b842ea133cf8d22cd7e34.tar.gz /tmp/liteadmin.tar.gz RUN mkdir /tmp/liteadmin && tar xzf /tmp/liteadmin.tar.gz --strip-components=1 -C /tmp/liteadmin && \ uv venv /opt/liteadmin --python python3.13 && \ uv pip install --python /opt/liteadmin/bin/python --require-hashes -r /tmp/liteadmin/requirements.txt && \ @@ -141,7 +142,7 @@ USER root RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent +RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app diff --git a/README.md b/README.md index 5fd84b529dc..4feb6e00ffa 100644 --- a/README.md +++ b/README.md @@ -616,15 +616,15 @@ The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws ### Verify Docker Image Signatures -All LiteLLM Docker images published to GHCR are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Every release is signed with the same key introduced in [commit `0112e53`](https://github.com/BerriAI/litellm/commit/0112e53046018d726492c814b3644b7d376029d0). +All LiteLLM Docker images published to GHCR are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Releases are signed with the key introduced in [commit `e82a3c6`](https://github.com/BerriAI/litellm/commit/e82a3c646908ba0a6ee0beddc4310d7a6cac4065). **Verify using the pinned commit hash (recommended):** -A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the original signing key: +A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the published signing key: ```bash cosign verify \ - --key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \ + --key https://raw.githubusercontent.com/BerriAI/litellm/e82a3c646908ba0a6ee0beddc4310d7a6cac4065/cosign.pub \ ghcr.io/berriai/litellm: ``` diff --git a/security.md b/SECURITY.md similarity index 98% rename from security.md rename to SECURITY.md index c73379a9c01..d3603448d46 100644 --- a/security.md +++ b/SECURITY.md @@ -79,7 +79,7 @@ To qualify for a bounty, reports must include clear reproduction steps, a reprod - Audit Logs with retention policy - Control Allowed IP Addresses that can access your Cloud LiteLLM Instance -For security inquiries, please contact us at support@berri.ai +For security inquiries, please contact us at security@berri.ai #### Supported data regions for LiteLLM Cloud diff --git a/backend/Dockerfile b/backend/Dockerfile index 84c90e16f67..1c22f914100 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -43,7 +43,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/ # instead of nodeenv downloading one whose dynamic deps may not be in Wolfi # (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -109,7 +109,7 @@ ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG} USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \ + apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -124,7 +124,12 @@ ENV HOME=/home/nonroot \ PYTHONUNBUFFERED=1 \ PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries -COPY --from=builder --chown=nonroot:nonroot /app /app +COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv +COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm +COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise +COPY --from=builder --chown=nonroot:nonroot /app/backend /app/backend +COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh +COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma COPY --from=builder /opt/prisma /opt/prisma RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index 932da84e5d7..c1aeaabc5d2 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -4,8 +4,7 @@ Every pull request whose diff against its merge base touches one of the two cost 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. +get the file checks and additionally may only touch those two files and may only add or update models. """ from __future__ import annotations @@ -21,7 +20,6 @@ from typing import Final COST_MAP_PATH: Final = "model_prices_and_context_window.json" BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" -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"}) @@ -110,7 +108,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 and path != LEGACY_SCHEMA_PATH + if path not in GUARDED_PATHS ), *(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/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md new file mode 100644 index 00000000000..c79632e9195 --- /dev/null +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -0,0 +1,76 @@ +# Use your ChatGPT subscription through LiteLLM + +Sign into ChatGPT in Codex or Pi and send model requests through your existing LiteLLM proxy. The client keeps your login and refreshes it. Each request uses that client's ChatGPT account, while a separate LiteLLM virtual key controls gateway access and usage attribution. + +An administrator adds a `chatgpt/` model using the usual model configuration. No ChatGPT login needs to be added to the proxy. This uses the same separate-key convention as provider passthrough, including Claude Code subscription forwarding. + +## Add a model to the proxy + +Reuse your existing proxy. If you need one, follow the [proxy quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start). Merge this entry into its model list and restart the proxy: + +```yaml +model_list: + - model_name: gpt-6-luna + litellm_params: + model: chatgpt/gpt-6-luna + model_info: + mode: responses +``` + +Use a model available to the caller's subscription. The virtual key must allow the exposed model name. Leave saved ChatGPT logins and API credentials unset for a client-only deployment. Explicit cross-provider fallbacks can incur separate API charges; do not configure them for subscription-only traffic. + +## Configure Codex + +Use your existing LiteLLM URL and a virtual key assigned to your gateway user. Export that key as `LITELLM_API_KEY` in the environment that launches Codex. Keep the value out of committed configuration. + +Add this provider to `~/.codex/config.toml`, preserving your other settings: + +```toml +model = "gpt-6-luna" +model_provider = "litellm_subscription" + +[model_providers.litellm_subscription] +name = "LiteLLM subscription" +base_url = "https://litellm.example.com/v1" +wire_api = "responses" +supports_websockets = false +requires_openai_auth = true +env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" } +``` + +Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway. + +For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway). + +## Native compatibility and Pi + +For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name. + +### Configure Pi + +Sign into the built-in `openai-codex` provider in Pi. Merge this override into `~/.pi/agent/models.json` and export your LiteLLM virtual key as `LITELLM_API_KEY` before launching Pi: + +```json +{ + "providers": { + "openai-codex": { + "baseUrl": "https://litellm.example.com/backend-api/codex", + "headers": { + "x-litellm-api-key": "LITELLM_API_KEY" + } + } + } +} +``` + +Keep Pi's subscription login as the provider credential. The header value names an environment variable, not a key to paste into the file. + +## Verify account usage + +Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person. + +The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions. + +The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate. + +This flow does not connect accounts through the dashboard or pool subscription allowances. Use a Codex build that supports ChatGPT authentication with a custom model provider. The integration forwards HTTP and SSE requests; it does not add native WebSocket or other Codex backend routes. diff --git a/cosign.pub b/cosign.pub index 2c2a555ab09..7603d69641f 100644 --- a/cosign.pub +++ b/cosign.pub @@ -1,4 +1,4 @@ -----BEGIN PUBLIC KEY----- -MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEKi4ivqGpE231OGH50PKbqy1Y1Kkb -POJC8+i2Wko82gBOUCe3M0Vw86H/4rhUhfoYEti4gdJ9wZbYmK0I2EE96g== +MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEG06dpKM+vEDnmp2zcjyXgVCXxzRg +JA5H2LnmwBKQafZGka6OOrWujCEWATQewiAgyD4WQdJ0pvU9RC4mvNAjvQ== -----END PUBLIC KEY----- diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 56210458532..20f810686a3 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -59,6 +59,7 @@ RUN apk add --no-cache \ python-3.13 \ python-3.13-dev \ openssl \ + libssl3 \ openssl-dev \ nodejs \ npm \ @@ -122,7 +123,7 @@ ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG} USER root # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent +RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 086c38e0d60..90bcc523d9e 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,8 +1,8 @@ # syntax=docker/dockerfile:1.7 # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -60,6 +60,7 @@ RUN for i in 1 2 3; do \ coreutils \ curl \ openssl \ + libssl3 \ libsndfile \ nodejs \ npm && break || sleep 5; \ @@ -134,7 +135,7 @@ RUN for i in 1 2 3; do \ apk upgrade --no-cache && break || sleep 5; \ done && \ for i in 1 2 3; do \ - apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \ + apk add --no-cache python-3.13 bash openssl libssl3 tzdata libsndfile nodejs libevent && break || sleep 5; \ done COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer diff --git a/docker/build_from_pip/Dockerfile.build_from_pip b/docker/build_from_pip/Dockerfile.build_from_pip index a5733f0e1a0..25a7ed2fa97 100644 --- a/docker/build_from_pip/Dockerfile.build_from_pip +++ b/docker/build_from_pip/Dockerfile.build_from_pip @@ -30,7 +30,7 @@ RUN uv venv --python python && \ "google-cloud-aiplatform==1.133.0" \ "google-genai==1.37.0" \ "anthropic[vertex]==0.84.0" \ - "grpcio==1.78.0" \ + "grpcio==1.84.1" \ "prometheus-client==0.20.0" \ "langfuse==2.59.7" \ "opentelemetry-api==1.28.0" \ diff --git a/enterprise/enterprise_hooks/aporia_ai.py b/enterprise/enterprise_hooks/aporia_ai.py index 28b49bfce21..1e840f8dcc5 100644 --- a/enterprise/enterprise_hooks/aporia_ai.py +++ b/enterprise/enterprise_hooks/aporia_ai.py @@ -15,9 +15,10 @@ sys.path.insert( ) # Adds the parent directory to the system path import json import sys -from typing import Any, List, Literal, Optional +from typing import Final, List, Literal, Optional from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -29,11 +30,13 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.llms.openai.common_utils import decoded_json_object from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata from litellm.types.guardrails import GuardrailEventHooks GUARDRAIL_NAME = "aporia" +_DECODED_JSON: Final = TypeAdapter(object) class AporiaGuardrail(CustomGuardrail): @@ -68,7 +71,7 @@ class AporiaGuardrail(CustomGuardrail): async def prepare_aporia_request( self, new_messages: List[dict], response_string: Optional[str] = None ) -> dict: - data: dict[str, Any] = {} + data: dict[str, object] = {} if new_messages is not None: data["messages"] = new_messages if response_string is not None: @@ -122,8 +125,10 @@ class AporiaGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Aporia AI response: %s", response.text) if response.status_code == 200: # check if the response was flagged - _json_response = response.json() - action: str = _json_response.get( + _json_response: Final = decoded_json_object( + _DECODED_JSON.validate_python(response.json()) + ) + action: Final = _json_response.get( "action" ) # possible values are modify, passthrough, block, rephrase if action == "block": diff --git a/enterprise/enterprise_hooks/blocked_user_list.py b/enterprise/enterprise_hooks/blocked_user_list.py index f6998d58cd5..7e4d2514880 100644 --- a/enterprise/enterprise_hooks/blocked_user_list.py +++ b/enterprise/enterprise_hooks/blocked_user_list.py @@ -7,9 +7,10 @@ ## This accepts a list of user id's for whom calls will be rejected -from typing import Literal, Optional +from typing import Final, Literal, Optional from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._internal_context import with_service_target @@ -19,6 +20,9 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import LiteLLM_EndUserTable, UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import EndUserRepository + +_REQUEST_BODY_VALUE: Final = TypeAdapter(object) class ENTERPRISE_BlockedUserList(CustomLogger): @@ -94,15 +98,14 @@ class ENTERPRISE_BlockedUserList(CustomLogger): ) if end_user_cache_obj is None and self.prisma_client is not None: # check db - end_user_obj = ( - await self.prisma_client.db.litellm_endusertable.find_unique( - where={"user_id": user} - ) - ) + end_user_id: Final = _REQUEST_BODY_VALUE.validate_python(user) + end_user_obj = await EndUserRepository( + self.prisma_client + ).table.find_unique(where={"user_id": end_user_id}) if end_user_obj is None: # user not in db - assume not blocked end_user_obj = LiteLLM_EndUserTable(user_id=user, blocked=False) cache.set_cache(key=cache_key, value=end_user_obj, ttl=60) - if end_user_obj is not None and end_user_obj.blocked is True: + if end_user_obj.blocked is True: raise HTTPException( status_code=400, detail={ diff --git a/enterprise/enterprise_hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py index b3ef5a3d496..7dbbc9ed6ce 100644 --- a/enterprise/enterprise_hooks/google_text_moderation.py +++ b/enterprise/enterprise_hooks/google_text_moderation.py @@ -6,7 +6,10 @@ # +-----------------------------------------------+ # Thank you users! We ❤️ you! - Krrish & Ishaan +from typing import Final + from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -15,6 +18,8 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import iter_message_text from litellm.types.utils import CallTypesLiteral +_THRESHOLD_VALUE: Final = TypeAdapter(object) + class ENTERPRISE_GoogleTextModeration(CustomLogger): user_api_key_cache = None @@ -69,9 +74,11 @@ class ENTERPRISE_GoogleTextModeration(CustomLogger): f"{category}_confidence_threshold", default_confidence_threshold, ) - set_confidence_value = getattr( - self, - f"{category}_confidence_threshold", + set_confidence_value = _THRESHOLD_VALUE.validate_python( + getattr( + self, + f"{category}_confidence_threshold", + ) ) verbose_proxy_logger.info( f"Google Text Moderation: {category}_confidence_threshold: {set_confidence_value}" diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py b/enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py index 14d34f5d1e8..e0190ea8da0 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py @@ -1,15 +1,20 @@ # this is an example endpoint to receive data from litellm +from typing import Final + from fastapi import FastAPI, HTTPException, Request +from pydantic import TypeAdapter app = FastAPI() +_LOGGED_EVENT: Final = TypeAdapter(object) + @app.post("/log-event") async def log_event(request: Request): try: print("Received /log-event request") # noqa # Assuming the incoming request has JSON data - data = await request.json() + data: Final = _LOGGED_EVENT.validate_python(await request.json()) print("Received request data:") # noqa print(data) # noqa diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index 3422e8969b0..e12539eac53 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -8,25 +8,30 @@ ## This provides an LLM Guard Integration for content moderation on the proxy import asyncio -from typing import Final, Optional +from collections.abc import Mapping +from typing import Final import aiohttp from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.openai.common_utils import decoded_json_object from litellm.proxy._types import UserAPIKeyAuth from litellm.secret_managers.main import get_secret_str from litellm.types.utils import CallTypesLiteral +_DECODED_JSON: Final = TypeAdapter(object) + class _ENTERPRISE_LLMGuard(CustomLogger): # Class variables or attributes def __init__( self, mock_testing: bool = False, - mock_redacted_text: Optional[dict] = None, + mock_redacted_text: Mapping[str, object] | None = None, ): self.mock_redacted_text = mock_redacted_text self.llm_guard_mode = litellm.llm_guard_mode @@ -66,7 +71,9 @@ class _ENTERPRISE_LLMGuard(CustomLogger): async with session.post( analyze_url, json={"prompt": text} ) as response: - redacted_text = await response.json() + redacted_text = _DECODED_JSON.validate_python( + await response.json() + ) verbose_proxy_logger.debug( f"LLM Guard: Received response - {redacted_text}" ) @@ -77,12 +84,13 @@ class _ENTERPRISE_LLMGuard(CustomLogger): "error": f"Invalid content moderation response: {redacted_text}" }, ) - if redacted_text.get("is_valid", None) is False: + redacted_response: Final = decoded_json_object(redacted_text) + if redacted_response.get("is_valid", None) is False: raise HTTPException( status_code=400, detail={"error": "Violated content safety policy"}, ) - sanitized_prompt = redacted_text.get("sanitized_prompt") + sanitized_prompt: Final = redacted_response.get("sanitized_prompt") return sanitized_prompt if isinstance(sanitized_prompt, str) else text except Exception as e: verbose_proxy_logger.exception( diff --git a/enterprise/litellm_enterprise/proxy/auth/route_checks.py b/enterprise/litellm_enterprise/proxy/auth/route_checks.py index fc57292a8d2..479537c4a90 100644 --- a/enterprise/litellm_enterprise/proxy/auth/route_checks.py +++ b/enterprise/litellm_enterprise/proxy/auth/route_checks.py @@ -2,6 +2,8 @@ import os from fastapi import HTTPException, status +from litellm.proxy._types import MODEL_LIST_ROUTES + class EnterpriseRouteChecks: @staticmethod @@ -43,7 +45,7 @@ class EnterpriseRouteChecks: # Routes that should remain accessible even when LLM API endpoints are disabled. # These are read-only model listing routes needed by the Admin UI. - LLM_API_EXEMPT_ROUTES = ["/models", "/v1/models"] + LLM_API_EXEMPT_ROUTES = MODEL_LIST_ROUTES @staticmethod def should_call_route(route: str): diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index df1f263a947..f3b1240f733 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -1040,8 +1040,8 @@ class CheckBatchCost: from litellm.proxy.proxy_server import proxy_config from litellm.secret_managers.main import str_to_bool - api_key: Final = getattr(job, "api_key", None) - team_id: Final = getattr(job, "team_id", None) + api_key: Final = job.api_key if hasattr(job, "api_key") else None + team_id: Final = job.team_id if hasattr(job, "team_id") else None has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key) try: key_row: Final = ( diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bc35c4a0aba..48cbc0ab8ff 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -44,6 +44,7 @@ from litellm.constants import ( ) from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.internal_call_metadata import requests_background from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -133,7 +134,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient - Span = Union[_Span, Any] + Span = _Span InternalUsageCache = _InternalUsageCache PrismaClient = _PrismaClient else: @@ -490,6 +491,14 @@ _PROVIDER_FILE_ID_KEYS: Final = tuple(dict.fromkeys(chain.from_iterable(_PROVIDE _PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0) +class _SupportsKeyedGet(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +def _keyed_get(source: object, key: str, default: object = None) -> object: + return require_attributes(source, _SupportsKeyedGet, "get").get(key, default) + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -1611,7 +1620,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return holders if holders or "input_file_id" in request else healthy_deployments async def async_pre_call_deployment_hook( - self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] + self, kwargs: dict[str, object], call_type: Optional[CallTypes] ) -> Optional[dict]: """ Allow modifying the request just before it's sent to the deployment. @@ -1623,11 +1632,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) # model_info may be at top-level or nested under litellm_metadata # (batch/file operations use litellm_metadata) - model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) + model_id = cast(Optional[str], _keyed_get(kwargs.get("model_info", {}), "id")) if model_id is None: model_id = cast( Optional[str], - kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), + _keyed_get(_keyed_get(kwargs.get("litellm_metadata", {}), "model_info", {}), "id"), ) for accessor_key in accessor_keys: input_file_id = cast(Optional[str], kwargs.get(accessor_key)) @@ -1824,7 +1833,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ), ) - async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span) -> dict: + async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span | None) -> dict: """ Get model-specific file IDs for a list of proxy file IDs. Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id @@ -1866,7 +1875,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): llm_router: Optional[Router], _create_file_request: CreateFileRequest, target_model_names_list: List[str], - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, ) -> List[OpenAIFileObject]: if llm_router is None: raise Exception("LLM Router not initialized. Ensure models added to proxy.") @@ -1882,7 +1891,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): create_file_request: CreateFileRequest, llm_router: Router, target_model_names_list: List[str], - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, user_api_key_dict: UserAPIKeyAuth, ) -> OpenAIFileObject: responses = await self.create_file_for_each_model( @@ -1949,7 +1958,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_objects: List[OpenAIFileObject], create_file_request: CreateFileRequest, internal_usage_cache: InternalUsageCache, - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, target_model_names_list: List[str], ) -> OpenAIFileObject: ## GET THE FILE TYPE FROM THE CREATE FILE REQUEST @@ -2652,7 +2661,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Continue with other files even if one fails continue - def _get_content_type_from_file_object(self, file_object: Optional[Any]) -> str: + def _get_content_type_from_file_object(self, file_object: object) -> str: """ Determine content type from file object. diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index 0e9657aba96..fcaf58099b9 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient - Span = Union[_Span, Any] + Span = _Span InternalUsageCache = _InternalUsageCache PrismaClient = _PrismaClient else: @@ -107,7 +107,7 @@ class _PROXY_LiteLLMManagedVectorStores( llm_router: Router, model: str, request_data: dict[str, object] | VectorStoreCreateOptionalRequestParams, - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, ) -> VectorStoreCreateResponse: """ Create a vector store for a specific model. @@ -134,7 +134,7 @@ class _PROXY_LiteLLMManagedVectorStores( create_request: VectorStoreCreateOptionalRequestParams, llm_router: Router, target_model_names_list: List[str], - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, user_api_key_dict: UserAPIKeyAuth, ) -> VectorStoreCreateResponse: """ diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index db419fde1e6..b91b4a970e8 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -36,6 +36,7 @@ from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.table_repositories import PrismaTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository @@ -73,6 +74,16 @@ def _verification_token_table( return VerificationTokenRepository(prisma_client).table +class _WriterVerificationTokenRepository(PrismaTableRepository["prisma_models.LiteLLM_VerificationToken"]): + table_name = "litellm_verificationtoken" + + +def _writer_verification_token_table( + prisma_client: PrismaClient, +) -> TableActions["prisma_models.LiteLLM_VerificationToken"]: + return _WriterVerificationTokenRepository(prisma_client, use_writer=True).table + + def _budget_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: return BudgetRepository(prisma_client).table @@ -849,7 +860,7 @@ async def update_project( else None ) if current_project is not None and data.team_id != current_project.team_id: - mismatched_key_count: Final = await prisma_client.writer_db.litellm_verificationtoken.count( + mismatched_key_count: Final = await _writer_verification_token_table(prisma_client).count( where={ "project_id": data.project_id, "OR": [{"team_id": {"not": data.team_id}}, {"team_id": None}], diff --git a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py index 9b87f4f1a7f..19d4988452b 100644 --- a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py +++ b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py @@ -14,6 +14,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Final, List, Optional, Protocol from fastapi import APIRouter, Depends, HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -37,6 +38,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient router = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) class ManagedVectorStoreRow(Protocol): @@ -348,11 +350,11 @@ async def update_vector_store( raise HTTPException(status_code=500, detail="Database not connected") try: - update_data = data.model_dump(exclude_unset=True) - vector_store_id: Final[str] = update_data.pop("vector_store_id") + update_data = data.model_dump(exclude_unset=True, exclude={"vector_store_id"}) + vector_store_id: Final = data.vector_store_id if update_data.get("vector_store_metadata") is not None: update_data["vector_store_metadata"] = safe_dumps( - update_data["vector_store_metadata"] + _DECODED_JSON.validate_python(update_data["vector_store_metadata"]) ) updated = await managed_vector_store_table(prisma_client).update( diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 051b1c30c8f..69068cfc8a5 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) ARG PGBOUNCER_VERSION=1.26.0 @@ -59,7 +59,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/ # instead of nodeenv downloading one whose dynamic deps may not be in Wolfi # (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -125,7 +125,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \ + apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic libevent && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -140,7 +140,12 @@ ENV HOME=/home/nonroot \ PYTHONUNBUFFERED=1 \ PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries -COPY --from=builder --chown=nonroot:nonroot /app /app +COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv +COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm +COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise +COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway +COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh +COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma COPY --from=builder /opt/prisma /opt/prisma COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 98f36607fff..cb50324e1dd 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -8202,18 +8202,18 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.9" +version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +checksum = "e49d0e65abc00c142e8f9428da991da9f1a250c39be97459901eb12b9b49f62a" dependencies = [ "rustls-pki-types", ] [[package]] name = "webpki-roots" -version = "1.0.9" +version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +checksum = "f5f280bbb2b287bc8e7de7e3852945a0ff50d1bbcb5ed42b6c7a568c546df171" dependencies = [ "rustls-pki-types", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index f04eb115071..9e5e0a96b73 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -124,7 +124,7 @@ moka = { version = "0.12.16", features = ["future"] } strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" percent-encoding = "2.3" -webpki-roots = "1" +webpki-roots = "1.0.10" time = { version = "0.3.53", features = ["parsing"] } criterion = "0.8.2" fancy-regex = "0.19.2" diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs index 162d325e4f4..8b63fadf5a6 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs @@ -51,7 +51,7 @@ mod rules; mod status; mod vertex_ai; -pub use original::{ExceptionFamily, OriginalException}; +pub use original::{ExceptionFamily, OriginalException, TimeoutPhase}; pub use public::{MappedFailure, PublicError, UpstreamResponse}; use rules::{Rule, contains_any, first_match}; @@ -67,7 +67,6 @@ const TIMEOUT_MARKERS: &[&str] = &[ pub struct ExceptionContext { pub model: String, pub custom_llm_provider: String, - pub asynchronous: bool, pub vertex_project: Option, pub vertex_location: Option, pub model_group: Option, @@ -105,11 +104,12 @@ pub fn exception_type( (None, message.clone(), None) } OriginalException::Timeout { + phase, timeout_seconds, elapsed_seconds, } => ( None, - timeout_message(context.asynchronous, *timeout_seconds, *elapsed_seconds), + timeout_message(*phase, *timeout_seconds, *elapsed_seconds), None, ), }; @@ -161,21 +161,16 @@ fn family_rules(family: ExceptionFamily) -> &'static [Rule] { } } -/// The text the Python HTTP handler's timeout carries: the sync and async handlers word it -/// differently. -fn timeout_message( - asynchronous: bool, +/// The text the Python HTTP handler's timeout carries: the phase that expired and the limit +/// httpx enforced for it, both rendered as Python would print them. +pub fn timeout_message( + phase: TimeoutPhase, timeout_seconds: Option, elapsed_seconds: Option, ) -> String { let timeout = python_float(timeout_seconds); - if asynchronous { - let elapsed = - python_float(elapsed_seconds.map(|seconds| (seconds * 1000.0).round() / 1000.0)); - format!("Connection timed out. Timeout passed={timeout}, time taken={elapsed} seconds") - } else { - format!("Connection timed out after {timeout} seconds.") - } + let elapsed = python_float(elapsed_seconds.map(|seconds| (seconds * 1000.0).round() / 1000.0)); + format!("{phase} timed out. Timeout passed={timeout}, time taken={elapsed} seconds") } fn python_float(value: Option) -> String { @@ -356,6 +351,7 @@ mod tests { #[test] fn a_handler_timeout_is_a_408_without_a_response() { let original = OriginalException::Timeout { + phase: TimeoutPhase::Read, timeout_seconds: Some(0.5), elapsed_seconds: Some(0.5031), }; @@ -363,7 +359,7 @@ mod tests { mapped("mistral", &original), MappedFailure { error: PublicError::Timeout { status: 408 }, - message: "MistralException - Connection timed out after 0.5 seconds.".into(), + message: "MistralException - Read timed out. Timeout passed=0.5, time taken=0.503 seconds".into(), upstream: None, debug_info: DEBUG.into(), } @@ -440,39 +436,44 @@ mod tests { } #[rstest::rstest] - #[case::sync( - false, + #[case::connect(TimeoutPhase::Connect, "Connection timed out.")] + #[case::read(TimeoutPhase::Read, "Read timed out.")] + #[case::write(TimeoutPhase::Write, "Write timed out.")] + #[case::pool(TimeoutPhase::Pool, "Connection pool timed out.")] + fn timeout_text_names_the_phase_that_expired( + #[case] phase: TimeoutPhase, + #[case] expected: &str, + ) { + assert_eq!( + timeout_message(phase, Some(0.5), Some(0.5031)), + format!("{expected} Timeout passed=0.5, time taken=0.503 seconds") + ); + } + + #[rstest::rstest] + #[case::rounds_the_elapsed_time( Some(0.5), Some(0.5031), - "Connection timed out after 0.5 seconds." - )] - #[case::async_rounds_the_elapsed_time( - true, - Some(0.5), - Some(0.5031), - "Connection timed out. Timeout passed=0.5, time taken=0.503 seconds" + "Timeout passed=0.5, time taken=0.503 seconds" )] #[case::whole_seconds_keep_a_decimal( - true, Some(600.0), Some(2.0), - "Connection timed out. Timeout passed=600.0, time taken=2.0 seconds" + "Timeout passed=600.0, time taken=2.0 seconds" )] #[case::unknown_values_render_as_none( - true, None, None, - "Connection timed out. Timeout passed=None, time taken=None seconds" + "Timeout passed=None, time taken=None seconds" )] - fn timeout_text_follows_the_delivery_mode( - #[case] asynchronous: bool, + fn timeout_text_renders_durations_as_python_floats( #[case] timeout_seconds: Option, #[case] elapsed_seconds: Option, #[case] expected: &str, ) { assert_eq!( - timeout_message(asynchronous, timeout_seconds, elapsed_seconds), - expected + timeout_message(TimeoutPhase::Read, timeout_seconds, elapsed_seconds), + format!("Read timed out. {expected}") ); } diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs index 4de67e6e7b5..6517b093bbd 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs @@ -10,6 +10,7 @@ pub enum OriginalException { message: String, }, Timeout { + phase: TimeoutPhase, timeout_seconds: Option, elapsed_seconds: Option, }, @@ -20,6 +21,19 @@ pub enum OriginalException { }, } +/// The httpx timeout that expired, spelled the way the Python HTTP handler names it. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::VariantArray)] +pub enum TimeoutPhase { + #[strum(serialize = "Connection")] + Connect, + #[strum(serialize = "Read")] + Read, + #[strum(serialize = "Write")] + Write, + #[strum(serialize = "Connection pool")] + Pool, +} + /// Which provider-specific text rules apply before the shared status table. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ExceptionFamily { diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index 1abea61b164..426b8b91182 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -72,26 +72,26 @@ fn duration(seconds: f64) -> Result { #[derive(Clone, Default)] pub(crate) struct CacheHeaders(std::sync::Arc>); -impl litellm_host::interceptors::Interceptors for CacheHeaders { +impl litellm_host::interceptors::Interceptors for CacheHeaders { async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, _: litellm_host::interceptors::RequestContext, - ) -> Result { + ) -> Result { Ok(wire) } async fn after_provider_response( &self, _: litellm_host::interceptors::RawResponse, - ) -> Result<(), litellm_inference::RouteError> { + ) -> Result<(), litellm_inference::Error> { Ok(()) } async fn result_ready( &self, facts: litellm_host::interceptors::ExecutionFacts, - ) -> Result<(), litellm_inference::RouteError> { + ) -> Result<(), litellm_inference::Error> { if let litellm_host::interceptors::ResultSource::Cache { key } = facts.source { let _ = self.0.set(key); } diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 39656b933eb..ec82a6ede12 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -4,7 +4,7 @@ use axum::{ response::{IntoResponse, Response}, }; use litellm_http::transport::Error as TransportError; -use litellm_inference::RouteError; +use litellm_inference::Error as InferenceError; use litellm_llms::base_llm::ocr::error::Error as OcrError; use serde_json::{Map, Value, json}; @@ -19,7 +19,7 @@ pub enum Error { )] UnknownModel(String), #[error(transparent)] - Route(#[from] RouteError), + Route(#[from] InferenceError), #[error(transparent)] Ocr(#[from] OcrError), #[error("{0} is not implemented by the Rust gateway")] @@ -36,8 +36,8 @@ impl IntoResponse for Error { } } -impl From> for Error { - fn from(error: litellm_host_http::Error) -> Self { +impl From> for Error { + fn from(error: litellm_host_http::Error) -> Self { match error { litellm_host_http::Error::Call(error) => Self::Route(error), litellm_host_http::Error::Protocol => Self::Internal(error.to_string()), @@ -50,7 +50,7 @@ impl Error { match self { Self::Auth(error) => error.status(), Self::Unsupported(_) - | Self::Route(RouteError::Unsupported(_)) + | Self::Route(InferenceError::Unsupported(_)) | Self::Ocr(OcrError::Unsupported(_)) => StatusCode::NOT_IMPLEMENTED, Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE, Self::Ocr( @@ -64,10 +64,10 @@ impl Error { .and_then(|status| StatusCode::from_u16(status).ok()) .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), Self::InvalidBody(_) | Self::UnknownModel(_) => StatusCode::BAD_REQUEST, - Self::Route(RouteError::Transport(TransportError::Http { status, .. })) => { + Self::Route(InferenceError::Transport(TransportError::Http { status, .. })) => { StatusCode::from_u16(*status).unwrap_or(StatusCode::BAD_GATEWAY) } - Self::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey { .. })) => { + Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => { StatusCode::UNAUTHORIZED } Self::Route(error) if error.is_request() => StatusCode::BAD_REQUEST, @@ -97,7 +97,9 @@ impl Error { /// body already in that shape passes through, any other has its message extracted. pub fn body(&self, request_id: Option<&str>) -> Value { let raw = match self { - Self::Route(RouteError::Transport(TransportError::Http { body, .. })) => body.clone(), + Self::Route(InferenceError::Transport(TransportError::Http { body, .. })) => { + body.clone() + } other => other.to_string(), }; let parsed = serde_json::from_str::(&raw).ok(); @@ -175,7 +177,7 @@ mod tests { use super::*; fn upstream(status: u16, body: &str) -> Error { - Error::Route(RouteError::Transport(TransportError::Http { + Error::Route(InferenceError::Transport(TransportError::Http { status, body: body.into(), })) @@ -225,16 +227,16 @@ mod tests { #[rstest] #[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)] - #[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] + #[case::rejected_request(Error::Route(InferenceError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] #[case::missing_key( - Error::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey { + Error::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", environment_variable: "ANTHROPIC_API_KEY", })), StatusCode::UNAUTHORIZED, )] #[case::lost_connection( - Error::Route(RouteError::Transport(TransportError::Network("reset".into()))), + Error::Route(InferenceError::Transport(TransportError::Network("reset".into()))), StatusCode::INTERNAL_SERVER_ERROR, )] fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { diff --git a/litellm-rust/crates/inference-chat/src/lib.rs b/litellm-rust/crates/inference-chat/src/lib.rs index f6b391e0ba1..3be83840901 100644 --- a/litellm-rust/crates/inference-chat/src/lib.rs +++ b/litellm-rust/crates/inference-chat/src/lib.rs @@ -1,7 +1,7 @@ use litellm_host::observation::ObservationSender; pub mod route; pub mod types; -pub use litellm_inference::RouteError as Error; +pub use litellm_inference::Error; mod common_utils; pub mod constants; pub(crate) mod handler; diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 3fc432f8b46..70da8bc3255 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -21,7 +21,7 @@ use litellm_host::{ lifecycle::{CallEvent, ExecutionEvent}, observation::observation_channel, }; -use litellm_inference::RouteError; +use litellm_inference::Error; use rstest::{fixture, rstest}; use serde_json::{Value, json}; @@ -166,12 +166,12 @@ struct ChangingHooks { facts: std::sync::Mutex>, } -impl Interceptors for ChangingHooks { +impl Interceptors for ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, _: RequestContext, - ) -> Result { + ) -> Result { let call = self.calls.fetch_add(1, Ordering::SeqCst); if self.rewrite { wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 }); @@ -179,11 +179,11 @@ impl Interceptors for ChangingHooks { Ok(wire) } - async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> { + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { Ok(()) } - async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.facts.lock().unwrap().push(facts); Ok(()) } diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index f91d70445f3..bfefb9c63b0 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -14,7 +14,7 @@ use std::sync::Arc; use crate::common_utils::messages_provider; -pub use litellm_inference::RouteError as Error; +pub use litellm_inference::Error; pub use types::{ MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, messages_body, diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 32c962091ad..480b3c4f3b7 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -17,7 +17,7 @@ use litellm_host::interceptors::{ ResultSource, WireRequest, }; use litellm_inference::{ - RouteError, + Error, caching::{CacheRequest, execute_unary}, }; use rstest::{fixture, rstest}; @@ -133,12 +133,12 @@ struct ChangingHooks { facts: std::sync::Mutex>, } -impl Interceptors for ChangingHooks { +impl Interceptors for ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, _: RequestContext, - ) -> Result { + ) -> Result { let call = self.calls.fetch_add(1, Ordering::SeqCst); if self.rewrite { wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 }); @@ -146,11 +146,11 @@ impl Interceptors for ChangingHooks { Ok(wire) } - async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> { + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { Ok(()) } - async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.facts.lock().unwrap().push(facts); Ok(()) } diff --git a/litellm-rust/crates/inference-ocr/AGENTS.md b/litellm-rust/crates/inference-ocr/AGENTS.md index d74b7e7af06..e08f561e75d 100644 --- a/litellm-rust/crates/inference-ocr/AGENTS.md +++ b/litellm-rust/crates/inference-ocr/AGENTS.md @@ -3,4 +3,4 @@ - The route returns its completed response directly - Route closures supply OCR-specific host capabilities, such as the caller's Azure AD token provider (`src/route.rs`) - Provider code reaches the caller's hooks mid-call only through `litellm_llms::base_llm::ocr::handler::CallHooks`, which this crate implements over its host until it folds into `litellm_host::interceptors::Interceptors` -- OCR failures use `base_llm/ocr/error.rs`, the recorded exception to `litellm_inference::RouteError` until OCR folds into `litellm_llms::Error` +- OCR failures use `base_llm/ocr/error.rs`, the recorded exception to `litellm_inference::Error` until OCR folds into `litellm_llms::Error` diff --git a/litellm-rust/crates/inference-ocr/src/provider_config.rs b/litellm-rust/crates/inference-ocr/src/provider_config.rs index 2f0346f6e47..09041e9b72b 100644 --- a/litellm-rust/crates/inference-ocr/src/provider_config.rs +++ b/litellm-rust/crates/inference-ocr/src/provider_config.rs @@ -249,7 +249,7 @@ pub(crate) fn resolve_provider_config<'a>( provider: LlmProviders::Mistral, } } - Err(litellm_inference::RouteError::InvalidProvider(provider)) => { + Err(litellm_inference::Error::InvalidProvider(provider)) => { return Err(Error::InvalidProvider(provider)); } Err(error) => return Err(Error::InvalidProvider(error.to_string())), diff --git a/litellm-rust/crates/inference-responses/src/lib.rs b/litellm-rust/crates/inference-responses/src/lib.rs index 6a9d5372311..b6e151ecec3 100644 --- a/litellm-rust/crates/inference-responses/src/lib.rs +++ b/litellm-rust/crates/inference-responses/src/lib.rs @@ -1,5 +1,5 @@ use litellm_host::observation::ObservationSender; -pub use litellm_inference::RouteError as Error; +pub use litellm_inference::Error; pub mod ids; pub mod route; diff --git a/litellm-rust/crates/inference-responses/tests/caching.rs b/litellm-rust/crates/inference-responses/tests/caching.rs index f45cefdc13b..462461695c3 100644 --- a/litellm-rust/crates/inference-responses/tests/caching.rs +++ b/litellm-rust/crates/inference-responses/tests/caching.rs @@ -18,7 +18,7 @@ use litellm_host::interceptors::{ WireRequest, }; use litellm_inference::{ - RouteError, + Error, caching::{CacheRequest, execute_unary}, }; use rstest::{fixture, rstest}; @@ -124,12 +124,12 @@ struct ChangingHooks { facts: std::sync::Mutex>, } -impl Interceptors for ChangingHooks { +impl Interceptors for ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, _: RequestContext, - ) -> Result { + ) -> Result { let call = self.calls.fetch_add(1, Ordering::SeqCst); if self.rewrite { wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 }); @@ -137,11 +137,11 @@ impl Interceptors for ChangingHooks { Ok(wire) } - async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> { + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { Ok(()) } - async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.facts.lock().unwrap().push(facts); Ok(()) } diff --git a/litellm-rust/crates/inference-transcription/src/lib.rs b/litellm-rust/crates/inference-transcription/src/lib.rs index 47e7bc70b31..f11e3abcd54 100644 --- a/litellm-rust/crates/inference-transcription/src/lib.rs +++ b/litellm-rust/crates/inference-transcription/src/lib.rs @@ -1,5 +1,5 @@ pub mod types; -pub use litellm_inference::RouteError as Error; +pub use litellm_inference::Error; mod constants; mod handler; diff --git a/litellm-rust/crates/inference/AGENTS.md b/litellm-rust/crates/inference/AGENTS.md index 2c4bd00108b..7f2bd0387be 100644 --- a/litellm-rust/crates/inference/AGENTS.md +++ b/litellm-rust/crates/inference/AGENTS.md @@ -2,7 +2,7 @@ - `gateway-inference` (serving) -> `inference-` (one API's call) -> `inference` (shared base) - This crate is the shared base every format crate builds on - - `RouteError` (`src/error.rs`), `CallOptions` (`src/lib.rs`), `CallContext` (`src/context.rs`) + - `Error` (`src/error.rs`), `CallOptions` (`src/lib.rs`), `CallContext` (`src/context.rs`) - diagnostic spans (`src/diagnostic.rs`), outbound send and signing (`src/outbound.rs`), provider resolution (`src/provider.rs`), `CoreResources` (`src/resources.rs`) - response caching (`src/caching.rs`): `Cachable`, `StreamCachable`, `CacheRequest`, `CallCache`, `execute_unary`, `execute_streaming`, stream capture - Shared test helpers belong in litellm-inference-testing, used only as a dev dependency @@ -49,14 +49,14 @@ - `litellm-inference-` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...) - Provider code never imports from `inference` or `inference-*` - Handlers belong in `inference-*` or `llms`, never in a host crate -- Import every item from its canonical path. Never re-export another crate's items or give an item a second public path. The only allowed re-exports are a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`) and each format crate's `pub use litellm_inference::RouteError as Error;`, which keeps the pre-split `::Error` name +- Import every item from its canonical path. Never re-export another crate's items or give an item a second public path. The only allowed re-exports are a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`) and each format crate's `pub use litellm_inference::Error;`, which keeps the pre-split `::Error` name ## Error placement - The workspace `Error definitions` rules shape each crate's error. This section decides which crate and module a failure belongs to - A failure is declared once, by the lowest crate that raises it - every crate above nests it unchanged (`#[error(transparent)] Auth(#[from] litellm_auth::Error)`) or maps it once at its boundary, as `src/error.rs` does for `litellm_llms::Error` - - `RouteError` collects route failures for every format crate and never re-declares a variant a lower crate raises + - `Error` collects route failures for every format crate and never re-declares a variant a lower crate raises - Scope follows the concept, not the first caller - an error type under `litellm-llms`'s `/` is private to that provider: no other provider and nothing in `base_llm` may import it - a failure two providers or two routes can hit (wire framing, stream event decoding, a malformed provider response) belongs to the crate that owns the concept: `litellm-framer` for framing, `litellm_llms::Error` for the transformation layer diff --git a/litellm-rust/crates/inference/src/caching.rs b/litellm-rust/crates/inference/src/caching.rs index 6dfda20ead4..09d03c06215 100644 --- a/litellm-rust/crates/inference/src/caching.rs +++ b/litellm-rust/crates/inference/src/caching.rs @@ -24,9 +24,9 @@ use serde::{Deserialize, Serialize, de::DeserializeOwned}; use serde_json::Value; use tokio_util::codec::Decoder; -use crate::RouteError; +use crate::Error; -pub trait Cachable: Protocol { +pub trait Cachable: Protocol { const SURFACE: &'static str; fn reusable(_response: &Self::Response) -> bool { @@ -145,15 +145,15 @@ pub async fn execute_unary( request: CacheRequest, cache: Option>, options: Option, - interceptors: &impl Interceptors, + interceptors: &impl Interceptors, observers: Option<&ObservationSender>, provider: F, -) -> Result +) -> Result where P: Cachable, P::Response: Serialize + DeserializeOwned, F: FnOnce() -> Fut, - Fut: Future>, + Fut: Future>, { let identity = request.identity.clone(); crate::diagnostic::provider(&identity.model, &identity.provider); @@ -191,15 +191,15 @@ pub async fn execute_streaming( request: CacheRequest, cache: Option>, options: Option, - interceptors: &impl Interceptors, + interceptors: &impl Interceptors, observers: Option<&ObservationSender>, provider: F, -) -> Result, RouteError> +) -> Result, Error> where P: StreamCachable, P::Response: Serialize + DeserializeOwned, F: FnOnce() -> Fut, - Fut: Future, RouteError>>, + Fut: Future, Error>>, { let identity = request.identity.clone(); crate::diagnostic::provider(&identity.model, &identity.provider); @@ -299,9 +299,9 @@ impl CallCache

{ } fn capture_stream( - chunks: futures_util::stream::BoxStream<'static, Result>, + chunks: futures_util::stream::BoxStream<'static, Result>, session: CacheSession, -) -> futures_util::stream::BoxStream<'static, Result> { +) -> futures_util::stream::BoxStream<'static, Result> { stream::try_unfold( (chunks, Some(Vec::::new()), session), |(mut chunks, captured, session)| async move { @@ -330,7 +330,7 @@ fn capture_stream( { session.store(entry).await; } - Ok::<_, RouteError>(None) + Ok::<_, Error>(None) } } }, @@ -369,9 +369,9 @@ fn successful_stream(text: &str, terminal: &str) -> bool { async fn publish( facts: ExecutionFacts, - interceptors: &impl Interceptors, + interceptors: &impl Interceptors, observers: Option<&ObservationSender>, -) -> Result<(), RouteError> { +) -> Result<(), Error> { if let Some(observers) = observers { observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady { facts: facts.clone(), diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index 6fcbb3051b4..655b768dde1 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -5,7 +5,7 @@ use litellm_host::{ observation::ObservationSender, }; -use crate::{CallOptions, RouteError}; +use crate::{CallOptions, Error}; pub struct CallContext<'a, I> { pub interceptors: &'a I, @@ -13,7 +13,7 @@ pub struct CallContext<'a, I> { pub cache: CachePolicy, } -impl<'a, I: Interceptors> CallContext<'a, I> { +impl<'a, I: Interceptors> CallContext<'a, I> { pub fn new(interceptors: &'a I, options: CallOptions) -> Self { Self { interceptors, @@ -22,7 +22,7 @@ impl<'a, I: Interceptors> CallContext<'a, I> { } } - pub async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { + pub async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { if let Some(observers) = &self.observers { observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady { facts: facts.clone(), @@ -31,7 +31,7 @@ impl<'a, I: Interceptors> CallContext<'a, I> { self.interceptors.result_ready(facts).await } - pub async fn response_received(&self, body: &str) -> Result<(), RouteError> { + pub async fn response_received(&self, body: &str) -> Result<(), Error> { let raw = RawResponse { body: body.to_owned(), }; @@ -43,6 +43,6 @@ impl<'a, I: Interceptors> CallContext<'a, I> { self.interceptors .after_provider_response(raw) .await - .map_err(RouteError::post_call) + .map_err(Error::post_call) } } diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index e5f02754f3d..895c8c55e01 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -11,7 +11,7 @@ use litellm_http::transport::Error as TransportError; use litellm_llms::Error as LlmError; #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum RouteError { +pub enum Error { #[error("expected {expected}, got {actual}")] InvalidType { expected: &'static str, @@ -38,10 +38,10 @@ pub enum RouteError { #[error(transparent)] Secret(#[from] SecretError), #[error("post-call hook failed: {0}")] - PostCallHook(#[source] Arc), + PostCallHook(#[source] Arc), } -impl From for RouteError { +impl From for Error { fn from(fault: litellm_host::machine::MachineFault) -> Self { use litellm_host::machine::MachineFault; Self::InvalidRequest(match fault { @@ -51,7 +51,7 @@ impl From for RouteError { } } -impl RouteError { +impl Error { pub fn post_call(error: Self) -> Self { Self::PostCallHook(Arc::new(error)) } @@ -76,7 +76,7 @@ impl RouteError { } } -impl From for RouteError { +impl From for Error { fn from(error: LlmError) -> Self { match error { LlmError::InvalidType { expected, actual } => Self::InvalidType { expected, actual }, @@ -99,7 +99,7 @@ impl SecretError { } } -impl From for RouteError { +impl From for Error { fn from(error: litellm_secrets::Error) -> Self { Self::Secret(SecretError(Arc::new(error))) } @@ -115,22 +115,22 @@ impl Eq for SecretError {} #[cfg(test)] mod tests { - use super::RouteError; + use super::Error; use litellm_llms::{Error as LlmError, ErrorDetail}; use rstest::rstest; #[test] fn a_missing_api_key_is_the_environment_not_the_request() { assert!( - !RouteError::Auth(litellm_auth::Error::MissingApiKey { + !Error::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", environment_variable: "ANTHROPIC_API_KEY", }) .is_request() ); - assert!(RouteError::Auth(litellm_auth::Error::InvalidHeader).is_request()); - assert!(RouteError::InvalidRequest("top_k".into()).is_request()); - assert!(!RouteError::InvalidResponse("bad json".into()).is_request()); + assert!(Error::Auth(litellm_auth::Error::InvalidHeader).is_request()); + assert!(Error::InvalidRequest("top_k".into()).is_request()); + assert!(!Error::InvalidResponse("bad json".into()).is_request()); } #[rstest] #[case::request(true)] @@ -139,7 +139,7 @@ mod tests { let source = serde_json::from_str::("{").unwrap_err(); let source_message = source.to_string(); let detail = ErrorDetail::invalid("test payload", source); - let error = RouteError::from(if request { + let error = Error::from(if request { LlmError::InvalidRequest(detail) } else { LlmError::InvalidResponse(detail) diff --git a/litellm-rust/crates/inference/src/lib.rs b/litellm-rust/crates/inference/src/lib.rs index a168802e7ec..80771e206ca 100644 --- a/litellm-rust/crates/inference/src/lib.rs +++ b/litellm-rust/crates/inference/src/lib.rs @@ -7,7 +7,7 @@ pub mod outbound; pub mod provider; pub mod resources; -pub use error::RouteError; +pub use error::Error; #[derive(Clone, Default)] pub struct CallOptions { diff --git a/litellm-rust/crates/inference/src/provider.rs b/litellm-rust/crates/inference/src/provider.rs index 7e1269db1cf..305aca85f97 100644 --- a/litellm-rust/crates/inference/src/provider.rs +++ b/litellm-rust/crates/inference/src/provider.rs @@ -2,7 +2,7 @@ use litellm_core_utils::get_llm_provider_logic::{ CustomLlmProvider, LlmProviders, get_custom_llm_provider, }; -use crate::error::RouteError as Error; +use crate::error::Error; #[derive(Debug)] pub struct ResolvedProvider<'a> { diff --git a/litellm-rust/crates/inference/tests/caching.rs b/litellm-rust/crates/inference/tests/caching.rs index 1bd53695c30..a0a6a757565 100644 --- a/litellm-rust/crates/inference/tests/caching.rs +++ b/litellm-rust/crates/inference/tests/caching.rs @@ -25,7 +25,7 @@ use litellm_host::{ protocol::Protocol, }; use litellm_inference::{ - RouteError, + Error, caching::{Cachable, CacheRequest, StreamCachable, execute_streaming, execute_unary}, }; use rstest::{fixture, rstest}; @@ -36,7 +36,7 @@ struct TestRoute; impl Protocol for TestRoute { type Request = Value; type Response = Value; - type Error = RouteError; + type Error = Error; type HostCall = Infallible; type Chunk = Bytes; type StreamHead = (); @@ -204,7 +204,7 @@ async fn streamed( .chunks(3) .map(|bytes| Ok(Bytes::copy_from_slice(bytes))) .collect::>(); - let ending = fail.then_some(Err(RouteError::Unsupported("test transport failure"))); + let ending = fail.then_some(Err(Error::Unsupported("test transport failure"))); Ok(CallOutput::Stream { head: (), chunks: stream::iter(chunks.into_iter().chain(ending)).boxed(), @@ -215,7 +215,7 @@ async fn streamed( .unwrap() } -async fn consume(output: OutputOf) -> Result, RouteError> { +async fn consume(output: OutputOf) -> Result, Error> { let CallOutput::Stream { chunks, .. } = output else { panic!("expected a stream"); }; @@ -315,7 +315,7 @@ async fn a_provider_failure_never_populates_the_cache(cache: Arc for CacheHitAccounting { - async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { +impl Interceptors for CacheHitAccounting { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.calls.fetch_add(1, Ordering::SeqCst); let ResultSource::Cache { key } = facts.source else { panic!("expected cache source") @@ -482,7 +482,7 @@ impl Interceptors for CacheHitAccounting { ); self.key.set(key).unwrap(); if self.reject { - return Err(RouteError::Unsupported("cache accounting rejected")); + return Err(Error::Unsupported("cache accounting rejected")); } Ok(()) } @@ -491,11 +491,11 @@ impl Interceptors for CacheHitAccounting { &self, wire: WireRequest, _: RequestContext, - ) -> Result { + ) -> Result { Ok(wire) } - async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> { + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { Ok(()) } } @@ -571,7 +571,7 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure( if reject { assert!(matches!( result, - Err(RouteError::Unsupported("cache accounting rejected")) + Err(Error::Unsupported("cache accounting rejected")) )); } else { assert_eq!(result.unwrap(), expected); @@ -590,7 +590,7 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure( impl Protocol for UnaryTestRoute { type Request = Value; type Response = Value; - type Error = RouteError; + type Error = Error; type HostCall = Infallible; type Chunk = Infallible; type StreamHead = Infallible; diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index b1509313b88..8429c98a88a 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -235,7 +235,10 @@ impl AnthropicBeta { BetaProvider::Databricks, ], Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse], - Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 => &[], + Self::ComputerUse20250124 + | Self::ComputerUse20251124 + | Self::Context1m20250807 + | Self::ThinkingBindingControls20260801 => &[], Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], Self::Effort20251124 => &[BetaProvider::VertexAi], Self::FastMode20260201 | Self::Oauth20250420 => &[ @@ -272,7 +275,6 @@ impl AnthropicBeta { BetaProvider::Databricks, ], Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi], - Self::ThinkingBindingControls20260801 => &[BetaProvider::AzureAi], Self::ThinkingDisplayUpdates20260818 => &[BetaProvider::Databricks], Self::ToolSearchTool20251019 => &[ BetaProvider::Anthropic, diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index dc63f0d6353..4dab0297e73 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -1,5 +1,5 @@ use litellm_http::transport::Error as TransportError; -use litellm_inference::RouteError; +use litellm_inference::Error; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, @@ -19,9 +19,9 @@ pyo3::create_exception!( "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." ); -pub(crate) fn route_error_to_pyerr(error: RouteError) -> PyErr { +pub(crate) fn route_error_to_pyerr(error: Error) -> PyErr { match error { - RouteError::Transport(TransportError::Http { status, body }) => { + Error::Transport(TransportError::Http { status, body }) => { RustUpstreamError::new_err((status, body)) } other => by_fault(other.is_request(), other.to_string()), @@ -42,13 +42,13 @@ mod tests { use super::*; #[rstest::rstest] - #[case::unsupported(RouteError::Unsupported("test capability"), true)] - #[case::invalid_provider(RouteError::InvalidProvider("unknown".into()), true)] - #[case::invalid_request(RouteError::InvalidRequest("empty messages".into()), true)] + #[case::unsupported(Error::Unsupported("test capability"), true)] + #[case::invalid_provider(Error::InvalidProvider("unknown".into()), true)] + #[case::invalid_request(Error::InvalidRequest("empty messages".into()), true)] #[case::connection(TransportError::Connect("unreachable".into()).into(), false)] #[case::network(TransportError::Network("timed out".into()).into(), false)] - #[case::invalid_response(RouteError::InvalidResponse("missing usage".into()), false)] - fn route_failures_are_terminal(#[case] error: RouteError, #[case] is_request: bool) { + #[case::invalid_response(Error::InvalidResponse("missing usage".into()), false)] + fn route_failures_are_terminal(#[case] error: Error, #[case] is_request: bool) { Python::initialize(); Python::attach(|py| { let failure = route_error_to_pyerr(error); @@ -86,14 +86,12 @@ mod tests { fn missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors() { Python::initialize(); Python::attach(|py| { - let missing = - route_error_to_pyerr(RouteError::Auth(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: "ANTHROPIC_API_KEY", - })); + let missing = route_error_to_pyerr(Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + })); assert!(missing.is_instance_of::(py)); - let invalid = - route_error_to_pyerr(RouteError::Auth(litellm_auth::Error::InvalidHeader)); + let invalid = route_error_to_pyerr(Error::Auth(litellm_auth::Error::InvalidHeader)); assert!(invalid.is_instance_of::(py)); }); } diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index 3b590a522fe..c9eb60e47ed 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -7,7 +7,7 @@ use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{call::HostedMachine, protocol::Protocol}; use litellm_host_python::{PythonBinding, PythonHostCalls, from_py}; use litellm_http::transport::Error as TransportError; -use litellm_inference::RouteError; +use litellm_inference::Error; use litellm_secrets::source::SecretSource; use pyo3::{ exceptions::PyValueError, @@ -119,14 +119,14 @@ impl InferenceHost { public_response(py, self.module, response) } - pub fn error(&self, py: Python<'_>, error: RouteError) -> PyResult { - if let RouteError::Secret(source) = &error + pub fn error(&self, py: Python<'_>, error: Error) -> PyResult { + if let Error::Secret(source) = &error && let Some(original) = crate::secrets::python_error(py, source.source_error()) { return Ok(original); } let native = match error { - RouteError::Transport(TransportError::Http { status, body }) => { + Error::Transport(TransportError::Http { status, body }) => { let error = RustUpstreamError::new_err((status, body)); error .value(py) @@ -144,7 +144,7 @@ impl InferenceHost { } pub(super) trait InferenceRoute: Sized + 'static { - type Protocol: Protocol; + type Protocol: Protocol; const OPERATION: LoggingOperation; const SYNC_CALL_TYPE: &'static str; const ASYNC_CALL_TYPE: &'static str; diff --git a/litellm-rust/crates/python-bridge/src/secrets/provider.rs b/litellm-rust/crates/python-bridge/src/secrets/provider.rs index 568a0cd2228..5f4cd193fcf 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/provider.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/provider.rs @@ -88,20 +88,23 @@ fn aws_context( } fn read_timeout(value: Option<&Bound<'_, PyAny>>) -> PyResult> { - let Some(value) = value.filter(|value| !value.is_none()) else { - return Ok(None); - }; - let seconds = match value.extract::() { - Ok(value) => Some(value), - Err(_) => value.getattr("read")?.extract::>()?, - }; - seconds + read_timeout_seconds(value)? .map(|value| { Duration::try_from_secs_f64(value).map_err(|_| PyValueError::new_err("invalid timeout")) }) .transpose() } +pub(super) fn read_timeout_seconds(value: Option<&Bound<'_, PyAny>>) -> PyResult> { + let Some(value) = value.filter(|value| !value.is_none()) else { + return Ok(None); + }; + match value.extract::() { + Ok(value) => Ok(Some(value)), + Err(_) => value.getattr("read")?.extract::>(), + } +} + fn vault_context(params: Option<&Bound<'_, PyAny>>) -> PyResult { let params = params.and_then(|value| value.cast::().ok()); let nested = params diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault.rs b/litellm-rust/crates/python-bridge/src/secrets/vault.rs index 46b29839fe4..79281662848 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault.rs @@ -2,10 +2,12 @@ mod operation; pub(super) use operation::{Failure, FailureKind, FailureStage, delete, rotate, write}; +use litellm_core_utils::exception_mapping_utils::{TimeoutPhase, timeout_message}; use litellm_secrets::hashicorp::{Error, RawOperationError}; use pyo3::prelude::*; use super::mutation::{error_value, http_message, json_value}; +use super::provider::read_timeout_seconds; const HTTP_TRANSPORT_MODULE: &str = "litellm.rust_bridge.host.http_transport"; @@ -69,21 +71,17 @@ fn cause_message(py: Python<'_>, kind: FailureKind, context: &ErrorContext) -> P } FailureKind::Native(RawOperationError::Timeout { method, elapsed }) => { if method == "POST" { - let elapsed = py - .import("builtins")? - .call_method1("round", (elapsed.as_secs_f64(), 3))?; let kwargs = pyo3::types::PyDict::new(py); kwargs.set_item( "message", - format!( - "Connection timed out. Timeout passed={}, time taken={} seconds", - context.timeout.as_deref().unwrap_or("None"), - elapsed.str()? + timeout_message( + TimeoutPhase::Read, + context.timeout, + Some(elapsed.as_secs_f64()), ), )?; kwargs.set_item("model", "default-model-name")?; kwargs.set_item("llm_provider", "litellm-httpx-handler")?; - kwargs.set_item("headers", pyo3::types::PyDict::new(py))?; py.import("litellm.exceptions")? .getattr("Timeout")? .call((), Some(&kwargs))? @@ -159,7 +157,7 @@ fn response_text(py: Python<'_>, body: &[u8]) -> PyResult { #[derive(Default)] pub(super) struct ErrorContext { - timeout: Option, + timeout: Option, aiohttp: bool, } @@ -173,7 +171,7 @@ impl ErrorContext { return Ok(Self::default()); } Ok(Self { - timeout: timeout.map(|value| value.str()?.extract()).transpose()?, + timeout: read_timeout_seconds(timeout)?, aiohttp: py .import(HTTP_TRANSPORT_MODULE)? .getattr("uses_aiohttp_transport")? diff --git a/litellm/__init__.py b/litellm/__init__.py index 1784deeb797..c0141f43611 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -175,7 +175,7 @@ argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload. gcs_pub_sub_use_v1: Optional[bool] = False # if you want to use v1 gcs pubsub logged payload generic_api_use_v1: Optional[bool] = False # if you want to use v1 generic api logged payload -argilla_transformation_object: Optional[Dict[str, Any]] = None +argilla_transformation_object: Mapping[str, str] | None = None _async_input_callback: List[Union[str, Callable[..., object], "CustomLogger"]] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. @@ -333,7 +333,7 @@ in_memory_llm_clients_cache: "LLMClientCache" safe_memory_mode: bool = False enable_azure_ad_token_refresh: Optional[bool] = False # Proxy Authentication - auto-obtain/refresh OAuth2/JWT tokens for LiteLLM Proxy -proxy_auth: Optional[Any] = None +proxy_auth: _ProxyAuthHandlerType | None = None ### DEFAULT AZURE API VERSION ### AZURE_DEFAULT_API_VERSION = "2025-02-01-preview" # this is updated to the latest ### DEFAULT WATSONX API VERSION ### @@ -503,7 +503,7 @@ skill_search_embedding_model: Optional[str] = None # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) -public_model_groups_links: Dict[str, Union[str, Dict[str, Any]]] = {} +public_model_groups_links: dict[str, str | dict[str, object]] = {} #### REQUEST PRIORITIZATION ####### priority_reservation: Optional[Dict[str, Union[float, "PriorityReservationDict"]]] = None # priority_reservation_settings is lazy-loaded via __getattr__ @@ -682,6 +682,7 @@ darkbloom_models: Set = set() openinfer_models: Final[Set[str]] = set() v0_models: Set = set() morph_models: Set = set() +opencode_go_models: Final[set[str]] = set() # mutable-ok: populated when the shared model catalog loads lambda_ai_models: Set = set() inception_models: Set = set() hyperbolic_models: Set = set() @@ -970,6 +971,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje v0_models.add(key) elif value.get("litellm_provider") == "morph": morph_models.add(key) + elif value.get("litellm_provider") == "opencode_go": + opencode_go_models.add(key) elif value.get("litellm_provider") == "lambda_ai": lambda_ai_models.add(key) elif value.get("litellm_provider") == "inception": @@ -1137,6 +1140,7 @@ model_list = list( | openinfer_models | v0_models | morph_models + | opencode_go_models | lambda_ai_models | inception_models | black_forest_labs_models @@ -1254,6 +1258,7 @@ def _build_models_by_provider() -> dict: "openinfer": openinfer_models, "v0": v0_models, "morph": morph_models, + "opencode_go": opencode_go_models, "lambda_ai": lambda_ai_models, "inception": inception_models, "hyperbolic": hyperbolic_models, @@ -2289,6 +2294,10 @@ if TYPE_CHECKING: LiteLLM_UpperboundKeyGenerateParams, ) from litellm.utils import ModelResponseListIterator as _ModelResponseListIterator + from collections.abc import Awaitable + from litellm.types.utils import ModelResponse as _ModelResponseType + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper as _CustomStreamWrapperType + from litellm.proxy_auth import ProxyAuthHandler as _ProxyAuthHandlerType # Cost calculator functions cost_per_token: Callable[..., Tuple[float, float]] @@ -2312,7 +2321,7 @@ if TYPE_CHECKING: supports_audio_output: Callable[..., bool] supports_system_messages: Callable[..., bool] supports_reasoning: Callable[..., bool] - acreate: Callable[..., Any] + acreate: Callable[..., Awaitable[_ModelResponseType | _CustomStreamWrapperType]] get_max_tokens: Callable[..., int] get_model_info: Callable[..., _ModelInfoType] register_prompt_template: Callable[..., None] @@ -2610,8 +2619,10 @@ def __getattr__(name: str) -> Any: cancel_run, delete_run, ) + from types import MappingProxyType - return locals()[name] + evals_exports: Final[Mapping[str, object]] = MappingProxyType(locals()) + return evals_exports[name] raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 46126c444a6..1a7dc343745 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -19,7 +19,7 @@ import importlib import sys from collections.abc import Callable, Mapping from types import ModuleType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final from typing_extensions import ReadOnly, TypedDict @@ -456,7 +456,7 @@ def _lazy_import_http_handlers(name: str) -> object: params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"} # Create the client instance - provider_id: Final = cast(Any, "litellm_module_level_client") + provider_id: Final = "litellm_module_level_client" async_client: Final = get_async_httpx_client( llm_provider=provider_id, params=params, diff --git a/litellm/_uuid.py b/litellm/_uuid.py index e9578b7287f..55df233b553 100644 --- a/litellm/_uuid.py +++ b/litellm/_uuid.py @@ -1,10 +1,14 @@ """ Internal unified UUID helper. -Always uses fastuuid for performance. +Uses fastuuid where supported and the standard library on Python 3.15+. """ -import fastuuid as _uuid +import sys +import uuid as _uuid + +if sys.version_info < (3, 15): + import fastuuid as _uuid # Expose a module-like alias so callers can use: uuid.uuid4() uuid = _uuid diff --git a/litellm/a2a_protocol/cost_calculator.py b/litellm/a2a_protocol/cost_calculator.py index 31ca81c44dd..9a1431d6457 100644 --- a/litellm/a2a_protocol/cost_calculator.py +++ b/litellm/a2a_protocol/cost_calculator.py @@ -5,7 +5,10 @@ Supports dynamic cost parameters that allow platform owners to define custom costs per agent query or per token. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, SupportsFloat, SupportsIndex + +from typing_extensions import Buffer if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( @@ -15,6 +18,15 @@ else: LitellmLoggingObject = Any +def _custom_response_cost(model_call_details: Mapping[str, object]) -> float | None: + response_cost: Final = model_call_details.get("response_cost", None) + if response_cost is None: + return None + if isinstance(response_cost, (str, Buffer, SupportsFloat, SupportsIndex)): + return float(response_cost) + raise TypeError(f"float() argument must be a string or a real number, not {type(response_cost).__name__!r}") + + class A2ACostCalculator: @staticmethod def calculate_a2a_cost( @@ -45,9 +57,9 @@ class A2ACostCalculator: model_call_details: Final = litellm_logging_obj.model_call_details # Check if user set a custom response cost (backward compatibility) - response_cost: Final = model_call_details.get("response_cost", None) - if response_cost is not None: - return float(response_cost) + custom_response_cost: Final = _custom_response_cost(model_call_details) + if custom_response_cost is not None: + return custom_response_cost # Get litellm_params for cost parameters litellm_params: Final = model_call_details.get("litellm_params", {}) or {} diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 404bcca712f..1cea0194715 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -84,14 +84,23 @@ def _filter_reserved_headers( return filtered or None +def _calling_key_hash(litellm_params: Mapping[str, object]) -> str | None: + principal: Final = litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) + if principal is None or isinstance(principal, str): + return principal + if not principal: + return None + raise AttributeError(f"{type(principal).__name__!r} object has no attribute 'encode'", name="encode", obj=principal) + + def _request_scoped_runtime_session_id( params: Mapping[str, object], - litellm_params: Mapping[str, Any], + litellm_params: Mapping[str, object], ) -> str | None: context_id: Final = get_session_id_from_a2a_params(params) if not isinstance(context_id, str) or not context_id: return None - return scope_session_to_principal(context_id, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)) + return scope_session_to_principal(context_id, _calling_key_hash(litellm_params)) def _validate_runtime_session_id(session_id: str, model: str) -> str: @@ -109,6 +118,13 @@ def _validate_runtime_session_id(session_id: str, model: str) -> str: ) +def _agentcore_model_name(litellm_params: Mapping[str, object]) -> str: + model: Final = litellm_params.get("model", "") + if isinstance(model, str): + return model + raise AttributeError(f"{type(model).__name__!r} object has no attribute 'startswith'", name="startswith", obj=model) + + class BedrockAgentCoreA2ATransformation: """ Request/response transformation for Bedrock AgentCore A2A agents. @@ -152,11 +168,8 @@ class BedrockAgentCoreA2ATransformation: """ # Extract model and strip the "bedrock/" prefix # "bedrock/agentcore/arn:aws:..." → "agentcore/arn:aws:..." - model: Final = litellm_params.get("model", "") - if model.startswith("bedrock/"): - agentcore_model = model[len("bedrock/") :] - else: - agentcore_model = model + model: Final = _agentcore_model_name(litellm_params) + agentcore_model: Final = model.removeprefix("bedrock/") # Build optional_params from litellm_params (everything except model and custom_llm_provider) optional_params: Final = {k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider")} diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py index 57c4a4677d0..c2d581321cd 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py @@ -8,11 +8,31 @@ WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model: """ import asyncio -from collections.abc import AsyncIterator, Mapping -from typing import Any, Final +from collections.abc import AsyncIterator, Iterable, Mapping, Sequence +from itertools import chain +from typing import Final from uuid import uuid4 from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items +from litellm.llms.openai.common_utils import decoded_json_object + + +def _a2a_text_part(part: object) -> Sequence[object]: + part_fields: Final = as_str_mapping(part) + if not isinstance(part_fields, dict): + return () + if part_fields.get("kind") in (None, "", "text") and part_fields.get("text"): + return (part_fields["text"],) + return () + + +def _join_with_spaces(texts: Sequence[object]) -> str: + strings: Final = [text for text in texts if isinstance(text, str)] + if len(strings) == len(texts): + return " ".join(strings) + first_bad: Final = next(position for position, text in enumerate(texts) if not isinstance(text, str)) + raise TypeError(f"sequence item {first_bad}: expected str instance, {type(texts[first_bad]).__name__} found") class WatsonxOrchestrateTransformation: @@ -29,22 +49,16 @@ class WatsonxOrchestrateTransformation: return f"{cp4d_host.rstrip('/')}/orchestrate/cpd/instances/{instance_id}" @staticmethod - def extract_text_from_a2a_params(params: dict[str, Any]) -> str: + def extract_text_from_a2a_params(params: Mapping[str, object]) -> str: """ Extract user message text from A2A MessageSendParams. A2A format: params.message.parts[*] where part.kind == "text" """ - message: Final = params.get("message", {}) - parts: Final = message.get("parts", []) - texts: Final = [] - for part in parts: - if not isinstance(part, dict): - continue - kind = part.get("kind") - if kind in (None, "", "text") and part.get("text"): - texts.append(part["text"]) - return " ".join(texts) or "" + message: Final = decoded_json_object(params.get("message", {})) + parts: Final[Iterable[object]] = iter_items(message.get("parts", [])) + texts: Final = list(chain.from_iterable(map(_a2a_text_part, parts))) + return _join_with_spaces(texts) or "" @staticmethod def build_wxo_run_body( diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index 7400844bf28..17d1b86a609 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -4,15 +4,20 @@ Utility functions for A2A protocol. import hashlib from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import has_attributes if TYPE_CHECKING: from a2a.types import SendMessageRequest, SendStreamingMessageRequest +class _SupportsModelDump(Protocol): + def model_dump(self, *, mode: str) -> object: ... + + class A2ARequestUtils: """Utility class for A2A request/response processing.""" @@ -74,7 +79,7 @@ class A2ARequestUtils: @staticmethod def get_input_message_from_request( request: "SendMessageRequest | SendStreamingMessageRequest", - ) -> Any: + ) -> object: """ Extract the input message from an A2A request. @@ -84,7 +89,7 @@ class A2ARequestUtils: Returns: The message object/dict or None """ - params: Final = getattr(request, "params", None) + params: Final[object] = getattr(request, "params", None) if params is None: return None return getattr(params, "message", None) @@ -127,10 +132,12 @@ class A2ARequestUtils: # the dict branch — request-side parts are a2a-sdk Part RootModels whose # kind/text live on part.root, which the object branch cannot read. This # mirrors how the response side already works (it operates on model_dump). - input_message = A2ARequestUtils.get_input_message_from_request(request) - if input_message is not None and hasattr(input_message, "model_dump"): - input_message = input_message.model_dump(mode="json") - input_text: Final = A2ARequestUtils.extract_text_from_message(input_message) + input_message: Final = A2ARequestUtils.get_input_message_from_request(request) + input_text: Final = A2ARequestUtils.extract_text_from_message( + input_message.model_dump(mode="json") + if has_attributes(input_message, _SupportsModelDump, "model_dump") + else input_message + ) prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text) # Count output tokens diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 236b18fe428..d68de6c9d4d 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -67,7 +67,7 @@ "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", "text_editor_20241022": null, "text_editor_20250124": null, - "thinking-binding-controls-2026-08-01": null, + "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", "thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18", "token-efficient-tools-2025-02-19": null, "web-fetch-2025-09-10": "web-fetch-2025-09-10", diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 96c38b69104..fec1d8e4465 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -36,6 +36,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.offline_mode import is_offline_mode # Cache for the loaded configuration _BETA_HEADERS_CONFIG: dict | None = None @@ -44,6 +45,7 @@ _BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_na BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]] _BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) +_DECODED_JSON: Final = TypeAdapter(object) class GetAnthropicBetaHeadersConfig: @@ -75,7 +77,7 @@ class GetAnthropicBetaHeadersConfig: } @staticmethod - def _check_is_valid_dict(fetched_config: dict) -> bool: + def _check_is_valid_dict(fetched_config: object) -> bool: """Check if fetched config is a non-empty dict with expected structure.""" if not isinstance(fetched_config, dict): verbose_logger.warning( @@ -109,7 +111,7 @@ class GetAnthropicBetaHeadersConfig: return True @classmethod - def validate_beta_headers_config(cls, fetched_config: dict) -> bool: + def validate_beta_headers_config(cls, fetched_config: object) -> bool: """ Validate the integrity of a fetched beta headers config. @@ -118,7 +120,7 @@ class GetAnthropicBetaHeadersConfig: return cls._check_is_valid_dict(fetched_config) @staticmethod - def fetch_remote_beta_headers_config(url: str, timeout: int = 5) -> dict: + def fetch_remote_beta_headers_config(url: str, timeout: int = 5) -> object: """ Fetch the beta headers config from a remote URL. @@ -127,7 +129,7 @@ class GetAnthropicBetaHeadersConfig: """ response: Final = httpx.get(url, timeout=timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) def get_beta_headers_config(url: str) -> dict: @@ -145,7 +147,7 @@ def get_beta_headers_config(url: str) -> dict: Dict containing the beta headers configuration """ # Check if local-only mode is enabled - if os.getenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "").lower() == "true": + if is_offline_mode() or os.getenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "").lower() == "true": # verbose_logger.debug("Using local Anthropic beta headers config (LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS=True)") return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() @@ -160,7 +162,9 @@ def get_beta_headers_config(url: str) -> dict: return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() # Validate the fetched config - if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(fetched_config=content): + if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(fetched_config=content) or not isinstance( + content, dict + ): verbose_logger.warning( "LiteLLM: Fetched beta headers config failed integrity check. Using local backup instead. url=%s", url, diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index eb3ec8aaee2..d8198e41239 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -7,10 +7,14 @@ Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthrop import json from typing import Final +from pydantic import TypeAdapter + from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from .exceptions import AnthropicErrorDetail, AnthropicErrorResponse, AnthropicErrorType +_PARSED_ERROR_BODY: Final = TypeAdapter(object) + # HTTP status code -> Anthropic error type # Source: https://docs.anthropic.com/en/api/errors ANTHROPIC_ERROR_TYPE_MAP: Final[dict[int, AnthropicErrorType]] = { @@ -79,7 +83,7 @@ class AnthropicExceptionMapping: - Generic: {"message": "..."} - Plain strings """ - parsed: Final = safe_json_loads(raw_message) + parsed: Final = _PARSED_ERROR_BODY.validate_python(safe_json_loads(raw_message)) if isinstance(parsed, dict): # Bedrock format if "detail" in parsed and isinstance(parsed["detail"], dict): @@ -145,9 +149,8 @@ class AnthropicExceptionMapping: AnthropicErrorResponse dict """ # Try to parse as JSON once - parsed: dict | None = safe_json_loads(raw_message) - if not isinstance(parsed, dict): - parsed = None + decoded: Final = _PARSED_ERROR_BODY.validate_python(safe_json_loads(raw_message)) + parsed: Final[dict[str, object] | None] = decoded if isinstance(decoded, dict) else None # If parsed and already in Anthropic format - passthrough if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict(parsed): diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index a3428261636..a5721c6f047 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -77,7 +77,7 @@ def get_assistants( api_base: str | None = None, api_version: str | None = None, **kwargs, -) -> SyncCursorPage[Assistant]: +) -> SyncCursorPage[Assistant] | Coroutine[None, None, AsyncCursorPage[Assistant]]: aget_assistants: Final[bool | None] = kwargs.pop("aget_assistants", None) if aget_assistants is not None and not isinstance(aget_assistants, bool): raise Exception("Invalid value passed in for aget_assistants. Only bool or None allowed") @@ -100,7 +100,7 @@ def get_assistants( elif timeout is None: timeout = 600.0 - response: SyncCursorPage[Assistant] | None = None + response: SyncCursorPage[Assistant] | Coroutine[None, None, AsyncCursorPage[Assistant]] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -564,7 +564,7 @@ def create_thread( tool_resources: OpenAICreateThreadParamsToolResources | None = None, client: OpenAI | None = None, **kwargs, -) -> Thread: +) -> Thread | Coroutine[None, None, Thread]: """ - get the llm provider - if openai - route it there @@ -610,7 +610,7 @@ def create_thread( api_base: str | None = None api_key: str | None = None - response: Thread | None = None + response: Thread | Coroutine[None, None, Thread] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -735,7 +735,7 @@ def get_thread( thread_id: str, client=None, **kwargs, -) -> Thread: +) -> Thread | Coroutine[None, None, Thread]: """Get the thread object, given a thread_id""" aget_thread: Final = kwargs.pop("aget_thread", None) optional_params: Final = GenericLiteLLMParams(**kwargs) @@ -757,7 +757,7 @@ def get_thread( timeout = 600.0 api_base: str | None = None api_key: str | None = None - response: Thread | None = None + response: Thread | Coroutine[None, None, Thread] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -903,7 +903,7 @@ def add_message( metadata: dict | None = None, client=None, **kwargs, -) -> OpenAIMessage: +) -> OpenAIMessage | Coroutine[None, None, OpenAIMessage]: ### COMMON OBJECTS ### a_add_message: Final = kwargs.pop("a_add_message", None) _message_data: Final = MessageData(role=role, content=content, attachments=attachments, metadata=metadata) @@ -935,7 +935,7 @@ def add_message( timeout = 600.0 api_key: str | None = None api_base: str | None = None - response: OpenAIMessage | None = None + response: OpenAIMessage | Coroutine[None, None, OpenAIMessage] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -1065,7 +1065,7 @@ def get_messages( thread_id: str, client: Any | None = None, **kwargs, -) -> SyncCursorPage[OpenAIMessage]: +) -> SyncCursorPage[OpenAIMessage] | Coroutine[None, None, AsyncCursorPage[OpenAIMessage]]: aget_messages: Final = kwargs.pop("aget_messages", None) optional_params: Final = GenericLiteLLMParams(**kwargs) litellm_params_dict: Final = get_litellm_params(**kwargs) @@ -1086,7 +1086,7 @@ def get_messages( elif timeout is None: timeout = 600.0 - response: SyncCursorPage[OpenAIMessage] | None = None + response: SyncCursorPage[OpenAIMessage] | Coroutine[None, None, AsyncCursorPage[OpenAIMessage]] | None = None api_key: str | None = None api_base: str | None = None if custom_llm_provider == "openai": @@ -1256,7 +1256,7 @@ def run_thread( client: object | None = None, event_handler: AssistantEventHandler | None = None, # for stream=True calls **kwargs, -) -> Run: +) -> Run | Coroutine[None, None, Run]: """Run a given thread + assistant.""" arun_thread: Final = kwargs.pop("arun_thread", None) optional_params: Final = GenericLiteLLMParams(**kwargs) @@ -1278,7 +1278,7 @@ def run_thread( elif timeout is None: timeout = 600.0 - response: Run | None = None + response: Run | Coroutine[None, None, Run] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py index d82e58ecc91..0e7589d88ab 100644 --- a/litellm/caching/_embedding_router.py +++ b/litellm/caching/_embedding_router.py @@ -13,7 +13,7 @@ This module is dependency-injected: callers pass the proxy ``llm_router`` and from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import litellm from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS @@ -25,12 +25,12 @@ if TYPE_CHECKING: def resolve_embedding_router( embedding_model: str, llm_router: Router | None, - llm_model_list: list[dict[str, Any]] | None, + llm_model_list: Sequence[Mapping[str, object]] | None, ) -> Router | None: """Return ``llm_router`` iff it serves ``embedding_model`` as a deployment.""" if llm_router is None: return None - router_model_names: Final[list[str]] = ( + router_model_names: Final[Sequence[object]] = ( [m["model_name"] for m in llm_model_list if "model_name" in m] if llm_model_list is not None else [] ) if embedding_model in router_model_names: diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index 2d27574b098..d9a40c44d2b 100644 --- a/litellm/caching/affinity_cache.py +++ b/litellm/caching/affinity_cache.py @@ -24,6 +24,16 @@ if current == false then end if ARGV[3] then local decoded, stored = pcall(cjson.decode, current) + if ARGV[4] then + local legacy = current + if decoded and type(stored) == 'string' then legacy = stored end + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + if eligible['model_id'] == legacy then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end if decoded and type(stored) == 'table' then for _, eligible in ipairs(cjson.decode(ARGV[3])) do local matches = true @@ -68,6 +78,7 @@ def claim_affinity_pin_in_memory( ttl_seconds: int, *, eligible_values: tuple[Mapping[str, str], ...] | None = None, + preserve_legacy_model_ids: bool = False, ) -> object: """No await between read and write, so same-loop claims agree during a Redis outage.""" existing: Final[object] = cache.in_memory_cache.get_cache(cache_key) @@ -75,7 +86,12 @@ def claim_affinity_pin_in_memory( if _legacy_pin_matches(existing, pin_value): set_local_affinity_pin(cache, cache_key, pin_value, ttl_seconds) return existing - winner: Final = existing if existing is not None and existing in (eligible_values or ()) else pin_value + eligible_existing: Final = existing in (eligible_values or ()) or ( + preserve_legacy_model_ids + and isinstance(existing, str) + and any("model_id" in value and _legacy_pin_matches(existing, value) for value in (eligible_values or ())) + ) + winner: Final = existing if existing is not None and eligible_existing else pin_value set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) return winner @@ -94,6 +110,7 @@ async def claim_affinity_pin( ttl_seconds: int, *, eligible_values: tuple[Mapping[str, str], ...] | None = None, + preserve_legacy_model_ids: bool = False, ) -> object: """Return the authoritative first writer, replacing it only when it becomes ineligible. @@ -112,6 +129,7 @@ async def claim_affinity_pin( if eligible_values is not None else () ), + *(("1",) if preserve_legacy_model_ids and eligible_values is not None else ()), ) raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here object, await claim_script(keys=(cache_key,), args=args) @@ -124,4 +142,11 @@ async def claim_affinity_pin( return winner except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error) - return claim_affinity_pin_in_memory(cache, cache_key, pin_value, ttl_seconds, eligible_values=eligible_values) + return claim_affinity_pin_in_memory( + cache, + cache_key, + pin_value, + ttl_seconds, + eligible_values=eligible_values, + preserve_legacy_model_ids=preserve_legacy_model_ids, + ) diff --git a/litellm/caching/azure_blob_cache.py b/litellm/caching/azure_blob_cache.py index bf0e15421bc..5b9071dc920 100644 --- a/litellm/caching/azure_blob_cache.py +++ b/litellm/caching/azure_blob_cache.py @@ -13,10 +13,14 @@ import json from contextlib import suppress from typing import Final +from pydantic import TypeAdapter + from litellm._logging import print_verbose, verbose_logger from .base_cache import BaseCache +_DECODED_BLOB: Final = TypeAdapter(object) + class AzureBlobCache(BaseCache): def __init__(self, account_url, container) -> None: @@ -65,7 +69,7 @@ class AzureBlobCache(BaseCache): print_verbose(f"Get Azure Blob Cache: key: {key}") as_bytes: Final = self.container_client.download_blob(key).readall() as_str: Final = as_bytes.decode("utf-8") - cached_response: Final[object] = json.loads(as_str) + cached_response: Final = _DECODED_BLOB.validate_python(json.loads(as_str)) verbose_logger.debug( "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s", @@ -86,7 +90,7 @@ class AzureBlobCache(BaseCache): blob: Final = await self.async_container_client.download_blob(key) as_bytes: Final = await blob.readall() as_str: Final = as_bytes.decode("utf-8") - cached_response: Final[object] = json.loads(as_str) + cached_response: Final = _DECODED_BLOB.validate_python(json.loads(as_str)) verbose_logger.debug( "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s", key, diff --git a/litellm/caching/base_cache.py b/litellm/caching/base_cache.py index 51c169ba796..4dc894a64ba 100644 --- a/litellm/caching/base_cache.py +++ b/litellm/caching/base_cache.py @@ -9,14 +9,7 @@ Has 4 methods: """ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Final - -if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - - Span = _Span | Any -else: - Span = Any +from typing import Final class BaseCache(ABC): diff --git a/litellm/caching/disk_cache.py b/litellm/caching/disk_cache.py index 50939ad51ca..9470f23570a 100644 --- a/litellm/caching/disk_cache.py +++ b/litellm/caching/disk_cache.py @@ -1,15 +1,19 @@ import json from typing import TYPE_CHECKING, Any, Final +from pydantic import TypeAdapter + from .base_cache import BaseCache if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any +_DECODED_DISK_VALUE: Final = TypeAdapter(object) + class DiskCache(BaseCache): def __init__(self, disk_cache_dir: str | None = None): @@ -44,7 +48,7 @@ class DiskCache(BaseCache): original_cached_response: Final = self.disk_cache.get(key) if original_cached_response: try: - cached_response = json.loads(original_cached_response) + cached_response = _DECODED_DISK_VALUE.validate_python(json.loads(original_cached_response)) except Exception: cached_response = original_cached_response return cached_response diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index cb6426f64ff..1cc3e106aa0 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -95,7 +95,10 @@ class InMemoryCache(BaseCache): """ Check if a specific key is expired """ - return key in self.ttl_dict and self._clock() > self.ttl_dict[key] + try: + return self._clock() > self.ttl_dict[key] + except KeyError: + return False def _cache_keys(self) -> tuple[str, ...]: return tuple(self.cache_dict) # pyright: ignore[reportUnknownVariableType] # cache_dict is untyped legacy state @@ -227,7 +230,10 @@ class InMemoryCache(BaseCache): if key in self.cache_dict: if self.evict_element_if_expired(key): return None - original_cached_response: Final = self.cache_dict[key] + try: + original_cached_response: Final = self.cache_dict[key] + except KeyError: + return None try: cached_response = json.loads(original_cached_response) except Exception: diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index eb17facfb12..afe6f8a1eab 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1738,7 +1738,9 @@ class RedisCache(BaseCache): return key_value_dict @_redis_circuit_breaker_guard - async def async_get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): + async def async_get_cache( + self, key, parent_otel_span: Span | None = None, *, raise_on_error: bool = False, **kwargs + ): from redis.asyncio import Redis _redis_client: Final[Redis] = self.init_async_client() @@ -1781,6 +1783,8 @@ class RedisCache(BaseCache): ) ) print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e}") + if raise_on_error: + raise _record_swallowed_redis_failure(self._circuit_breaker, e) @_redis_circuit_breaker_guard diff --git a/litellm/constants.py b/litellm/constants.py index 1c1dfb9aa1c..d1daf2ee2d1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -156,6 +156,9 @@ DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLIN DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1) DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250) HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0 +HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( + "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS", 300, 1, 86400 +) # Maximum wall-clock seconds a streaming response is allowed to run. # Streams exceeding this duration are terminated with a Timeout error. @@ -477,6 +480,7 @@ REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_sp REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer" +REDIS_TOOL_DISCOVERY_BUFFER_KEY: Final = "litellm_tool_discovery_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = get_env_int("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100) REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer" REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000 @@ -608,6 +612,7 @@ NON_LLM_CONNECTION_TIMEOUT: Final = get_env_int( ) # timeout for adjacent services (e.g. jwt auth) MAX_EXCEPTION_MESSAGE_LENGTH: Final = get_env_int("MAX_EXCEPTION_MESSAGE_LENGTH", 2000) MAX_STRING_LENGTH_PROMPT_IN_DB: Final = get_env_int("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048) +MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB: Final = 64 * 1024 BEDROCK_MAX_POLICY_SIZE: Final = get_env_int("BEDROCK_MAX_POLICY_SIZE", 75) # One entry per distinct AWS credential-argument set. Per-user cost attribution passes the attributed # identity as aws_session_name, so this bounds how many attributed identities keep a cached STS session. @@ -857,6 +862,7 @@ LITELLM_CHAT_PROVIDERS: Final = [ "heroku", "oci", "morph", + "opencode_go", "lambda_ai", "inception", "vercel_ai_gateway", @@ -1112,6 +1118,7 @@ openai_compatible_providers: Final[list[str]] = [ "v0", "helicone", "morph", + "opencode_go", "lambda_ai", "inception", "hyperbolic", @@ -1838,6 +1845,7 @@ LITELLM_UI_SESSION_DURATION: Final = os.getenv("LITELLM_UI_SESSION_DURATION", "2 ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME: Final = "db_spend_update_job" +TOOL_DISCOVERY_COMMIT_JOB_NAME: Final = "tool_discovery_commit_job" DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME: Final = "db_daily_tag_spend_update_job" PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME: Final = "prometheus_emit_budget_metrics" CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f952883bc87..8331c68fe24 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( _SERVICE_TIER_TO_COST_KEY_SUFFIX, BilledTokenRates, CostCalculatorUtils, + calculate_cache_writing_cost, calculate_cost_component, generic_cost_per_character, generic_cost_per_token, @@ -1200,12 +1201,16 @@ def get_usage_object( elif isinstance(usage_obj, dict): return Usage(**usage_obj) elif isinstance(usage_obj, BaseModel): - return Usage(**usage_obj.model_dump()) + return _usage_from_model(usage_obj) else: verbose_logger.debug("Unknown usage object type: %s, usage_obj: %s", type(usage_obj), usage_obj) return None +def _usage_from_model(usage_model: BaseModel) -> Usage: + return Usage(**usage_model.model_dump()) + + def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper: match usage: case DecisionsUsage(): @@ -1546,6 +1551,8 @@ def completion_cost( """ try: call_type = infer_call_type(call_type, completion_response) or "completion" + if isinstance(completion_response, DecisionsResponse) and completion_response.usage is None: + raise ValueError("Cannot calculate Decisions cost: provider did not report usage") if pricing_provider is not None: custom_llm_provider = pricing_provider # rebind-ok: model_info.pricing_provider overrides the routing provider for cost-map lookups @@ -1673,8 +1680,10 @@ def completion_cost( verbose_logger.debug("selected model name for cost calculation: %s", model) if completion_response is not None and ( - isinstance(completion_response, BaseModel) or isinstance(completion_response, dict) - ): # tts returns a custom class + isinstance(completion_response, (BaseModel, dict)) + or isinstance(completion_response, HttpxBinaryResponseContent) + and completion_response.usage is not None + ): if isinstance(completion_response, dict): usage_obj: dict | Usage | None = completion_response.get("usage", {}) else: @@ -1684,7 +1693,7 @@ def completion_cost( setattr( completion_response, "usage", - litellm.Usage(**_usage_for_dump.model_dump()), + _usage_from_model(_usage_for_dump), ) if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)): _usage = _decisions_usage(usage_obj).model_dump() @@ -2262,6 +2271,8 @@ def response_cost_calculator( ### VERTEX LOCATION ### vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") region_name: str | None = None, + custom_cost_per_token: CostPerToken | None = None, + custom_cost_per_second: float | None = None, ) -> float: """ Returns @@ -2290,6 +2301,8 @@ def response_cost_calculator( optional_params=optional_params, custom_pricing=custom_pricing, base_model=base_model, + custom_cost_per_token=custom_cost_per_token, + custom_cost_per_second=custom_cost_per_second, prompt=prompt, standard_built_in_tools_params=standard_built_in_tools_params, litellm_model_name=litellm_model_name, @@ -2873,6 +2886,7 @@ def _batch_rate( model_info: ModelInfo, key: Literal[ "input_cost_per_audio_token_batches", + "cache_read_input_audio_token_cost_batches", "input_cost_per_image_token_batches", "input_cost_per_video_token_batches", ], @@ -2958,7 +2972,12 @@ def batch_cost_calculator( total_completion_cost = 0.0 if batch_rates.input is not None: total_prompt_cost = _batch_prompt_cost( - usage, model_info, batch_rates.input, batch_rates.cache_read, batch_rates.cache_creation + usage, + model_info, + batch_rates.input, + batch_rates.cache_read, + batch_rates.cache_creation, + batch_rates.cache_creation_rates_by_ttl, ) elif input_cost_per_token: details: Final = parse_prompt_tokens_details(usage) @@ -3003,10 +3022,13 @@ def _batch_prompt_cost( input_rate: float, cache_read_rate: float | None, cache_creation_rate: float | None, + cache_creation_rates_by_ttl: Mapping[str, float], ) -> float: details: Final = parse_prompt_tokens_details(usage) cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0 - written_tokens: Final = details["cache_creation_tokens"] if cache_creation_rate is not None else 0 + cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0 + bills_cache_writes: Final = cache_creation_rate is not None or bool(cache_creation_rates_by_ttl) + written_tokens: Final = details["cache_creation_tokens"] if bills_cache_writes else 0 audio_tokens, image_tokens, video_tokens = ( details["audio_tokens"], details["image_tokens"], @@ -3020,13 +3042,27 @@ def _batch_prompt_cost( + audio_tokens * _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate) + image_tokens * _batch_rate(model_info, "input_cost_per_image_token_batches", input_rate) + video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate) - + cached_tokens * (cache_read_rate or 0.0) - + written_tokens * (cache_creation_rate or 0.0) + + (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0) + + cached_audio_tokens + * _batch_rate(model_info, "cache_read_input_audio_token_cost_batches", cache_read_rate or 0.0) + + ( + calculate_cache_writing_cost( + written_tokens, + details["cache_creation_token_details"], + input_rate if cache_creation_rate is None else cache_creation_rate, + cache_creation_rates_by_ttl, + ) + if bills_cache_writes + else 0.0 + ) ) +_ATTRIBUTE_VALUE: Final = TypeAdapter(object) + + def _attribute_value(obj: object, name: str) -> object: - return getattr(obj, name) + return _ATTRIBUTE_VALUE.validate_python(getattr(obj, name)) def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]: @@ -3661,7 +3697,26 @@ def _get_model_info_or_none(model: str, custom_llm_provider: str) -> ModelInfo | return None -def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[str, ...]) -> float | None: +_TranscriptionRateKey: TypeAlias = Literal[ + "input_cost_per_second", "input_cost_per_audio_token", "input_cost_per_token", "output_cost_per_token" +] + + +def _model_info_rate(info: ModelInfo, key: _TranscriptionRateKey) -> float | None: + match key: + case "input_cost_per_second": + return info.get("input_cost_per_second") + case "input_cost_per_audio_token": + return info.get("input_cost_per_audio_token") + case "input_cost_per_token": + return info.get("input_cost_per_token") + case "output_cost_per_token": + return info.get("output_cost_per_token") + case _: + assert_never(key) + + +def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[_TranscriptionRateKey, ...]) -> float | None: """First of ``keys`` this entry prices, read off the raw ``litellm.model_cost`` entry because ``get_model_info`` synthesizes zero token rates for entries that omit them.""" if info is None: @@ -3670,12 +3725,18 @@ def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[str, ...]) if declared is None: return None return next( - (float(value) for key in keys if declared.get(key) is not None and (value := info.get(key)) is not None), + ( + float(value) + for key in keys + if declared.get(key) is not None and (value := _model_info_rate(info, key)) is not None + ), None, ) -def _transcription_rate(keys: tuple[str, ...], override: ModelInfo | None, base: ModelInfo | None) -> float: +def _transcription_rate( + keys: tuple[_TranscriptionRateKey, ...], override: ModelInfo | None, base: ModelInfo | None +) -> float: rates: Final = (_declared_transcription_rate(info, keys) for info in (override, base)) return next((rate for rate in rates if rate is not None), 0.0) diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index 00d112773f1..faf2daecf5f 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -122,6 +122,7 @@ class SpeechToCompletionBridgeHandler: if isinstance(result, ModelResponse): return self.transformation_handler.transform_response( model_response=result, + model=model, response_format=requested_response_format if isinstance(requested_response_format, str) else None, ) else: diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 9c1e5ae7a29..ad6ed269adf 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final, cast from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_param if TYPE_CHECKING: @@ -38,6 +39,11 @@ class ChatAudioParam(TypedDict): format: ReadOnly[NotRequired[str]] +class ChatSpeechConfigAudioParam(TypedDict): + speech_config: ReadOnly[Mapping[str, object]] + format: ReadOnly[str] + + class SpeechToCompletionBridgeTransformationHandler: def _validate_response_format( self, model: str, custom_llm_provider: str, optional_params: Mapping[str, object] @@ -49,6 +55,20 @@ class SpeechToCompletionBridgeTransformationHandler: return from litellm.exceptions import BadRequestError + if custom_llm_provider.startswith("vertex_ai"): + from litellm.llms.vertex_ai.text_to_speech.transformation import VertexAITextToSpeechConfig + + vertex_supported: Final = ", ".join( + sorted({*GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS, *VertexAITextToSpeechConfig.GEMINI_FORMAT_MAPPINGS}) + ) + raise BadRequestError( + message=( + f"Vertex AI Gemini TTS does not support response_format='{response_format}'." + f" Supported response formats: {vertex_supported}." + ), + model=model, + llm_provider=custom_llm_provider, + ) supported: Final = ", ".join(sorted(GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS)) raise BadRequestError( message=( @@ -75,8 +95,27 @@ class SpeechToCompletionBridgeTransformationHandler: return response_format if isinstance(response_format, str) else None def _chat_audio_param( - self, model: str, voice: str | Mapping[str, object] | None, optional_params: Mapping[str, object] - ) -> ChatAudioParam | None: + self, + model: str, + voice: str | Mapping[str, object] | None, + optional_params: Mapping[str, object], + drop_params: bool, + custom_llm_provider: str, + ) -> ChatAudioParam | ChatSpeechConfigAudioParam | None: + import litellm + + if isinstance(voice, Mapping) and self._is_gemini_tts_model(litellm.model_alias_map.get(model, model)): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + normalize_gemini_speech_config, + reject_unknown_gemini_tts_voice_keys, + ) + + reject_unknown_gemini_tts_voice_keys(model, voice, drop_params, custom_llm_provider) + speech_audio: Final[ChatSpeechConfigAudioParam] = { + "speech_config": normalize_gemini_speech_config(voice, drops_unsupported_params(drop_params)), + "format": GEMINI_TTS_CHAT_AUDIO_FORMAT, + } + return speech_audio if not isinstance(voice, str): return None audio_format: Final = self._chat_audio_format(model, optional_params) @@ -90,9 +129,9 @@ class SpeechToCompletionBridgeTransformationHandler: self, model: str, input: str, - voice: str | dict | None, - optional_params: dict, - litellm_params: dict, + voice: str | Mapping[str, object] | None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], headers: dict, litellm_logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, @@ -104,7 +143,13 @@ class SpeechToCompletionBridgeTransformationHandler: "messages": [user_message], "modalities": ["audio"], **self._chat_completion_params(optional_params), - "audio": self._chat_audio_param(model, voice, optional_params), + "audio": self._chat_audio_param( + model, + voice, + optional_params, + drops_unsupported_params(litellm_params.get("drop_params")), + custom_llm_provider, + ), **litellm_params, "headers": headers, "litellm_logging_obj": litellm_logging_obj, @@ -153,8 +198,9 @@ class SpeechToCompletionBridgeTransformationHandler: return wav_header + pcm_data def _is_gemini_tts_model(self, model: str) -> bool: - """Check if the model is a Gemini TTS model that returns PCM16 data.""" - return "gemini" in model.lower() and ("tts" in model.lower() or "preview-tts" in model.lower()) + from litellm.utils import is_gemini_tts_model + + return is_gemini_tts_model(model) def _gemini_tts_response_body(self, decoded_audio: bytes, response_format: str | None) -> tuple[bytes, str]: if response_format == GEMINI_TTS_RAW_RESPONSE_FORMAT: @@ -162,7 +208,7 @@ class SpeechToCompletionBridgeTransformationHandler: return self._convert_pcm16_to_wav(decoded_audio), "audio/wav" def transform_response( - self, model_response: "ModelResponse", response_format: str | None + self, model_response: "ModelResponse", model: str, response_format: str | None ) -> "HttpxBinaryResponseContent": import base64 @@ -176,10 +222,9 @@ class SpeechToCompletionBridgeTransformationHandler: raise ValueError("No audio part found in the response") decoded_audio: Final = base64.b64decode(audio_part.data) - model: Final = getattr(model_response, "model", "") content, content_type = ( self._gemini_tts_response_body(decoded_audio, response_format) - if self._is_gemini_tts_model(model) + if self._is_gemini_tts_model(model) or self._is_gemini_tts_model(model_response.model or "") else (decoded_audio, "audio/mpeg") ) response: Final = httpx.Response( diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 733f1aa72d5..b8e933bf046 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1273,7 +1273,7 @@ class MidStreamFallbackError(ServiceUnavailableError): self.message = _saved_message self.args = (_saved_message,) - original_headers: Final = getattr(original_exception, "litellm_response_headers", None) + original_headers: Final[object] = getattr(original_exception, "litellm_response_headers", None) if original_headers is not None: self.litellm_response_headers = original_headers diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 8015b065899..4f204b3629e 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -10,6 +10,7 @@ from mcp.types import Tool as MCPTool from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( @@ -20,6 +21,8 @@ from litellm.constants import ( from litellm.types.llms.anthropic import AnthropicMessagesTool from litellm.types.utils import ChatCompletionMessageToolCall +_DECODED_JSON: Final = TypeAdapter(object) + ######################################################## # List MCP Tool functions @@ -200,15 +203,18 @@ async def call_mcp_tool( return tool_result +def _decode_json_arguments(arguments: str) -> object: + try: + return _DECODED_JSON.validate_python(json.loads(arguments)) + except json.JSONDecodeError: + return {} + + def _get_function_arguments(function: FunctionDefinition) -> dict: """Helper to safely get and parse function arguments.""" - arguments = function.get("arguments", {}) - if isinstance(arguments, str): - try: - arguments = json.loads(arguments) - except json.JSONDecodeError: - arguments = {} - return arguments if isinstance(arguments, dict) else {} + arguments: Final = function.get("arguments", {}) + decoded: Final = _decode_json_arguments(arguments) if isinstance(arguments, str) else arguments + return decoded if isinstance(decoded, dict) else {} def transform_openai_tool_call_request_to_mcp_tool_call_request( diff --git a/litellm/files/main.py b/litellm/files/main.py index 16ff5b02f9c..c5172bc01ae 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -9,7 +9,7 @@ import asyncio import contextvars import time import uuid as uuid_module -from collections.abc import Coroutine +from collections.abc import Awaitable, Coroutine from functools import partial from typing import Final, Literal, cast @@ -293,6 +293,14 @@ def create_file( raise e +async def _await_retrieved_file(pending: Awaitable[FileObject | OpenAIFileObject]) -> FileObject | OpenAIFileObject: + return await pending + + +async def _await_deleted_file(pending: Awaitable[object]) -> object: + return await pending + + @client async def afile_retrieve( file_id: str, @@ -325,7 +333,7 @@ async def afile_retrieve( func_with_context: Final = partial(ctx.run, func) init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): - response = await init_response + response = await _await_retrieved_file(init_response) else: response = init_response @@ -505,7 +513,7 @@ async def afile_delete( func_with_context: Final = partial(ctx.run, func) init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): - response = await init_response + response = await _await_deleted_file(init_response) else: response = init_response diff --git a/litellm/files/streaming.py b/litellm/files/streaming.py index 460050872f4..e9a030e0e32 100644 --- a/litellm/files/streaming.py +++ b/litellm/files/streaming.py @@ -1,7 +1,7 @@ import datetime import traceback from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Final, Optional, cast +from typing import TYPE_CHECKING, Final, Optional, cast import anyio @@ -163,7 +163,7 @@ class FileContentStreamingResponse: merged_hidden_params: Final = cast( "StandardLoggingHiddenParams", { - **cast(dict[str, Any], payload.get("hidden_params") or {}), + **(payload.get("hidden_params") or {}), **self._hidden_params, }, ) diff --git a/litellm/harness/endpoint.py b/litellm/harness/endpoint.py index 0abf7c11280..71e2742fa2e 100644 --- a/litellm/harness/endpoint.py +++ b/litellm/harness/endpoint.py @@ -37,6 +37,7 @@ from litellm.harness.context import GatewayTarget from litellm.harness.errors import HarnessError, HarnessInstallFailed from litellm.harness.types import Harness, Usage from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -690,18 +691,21 @@ class ModelEndpoint: async def _invoke_sdk( self, route: str, - kwargs: dict[str, Any], # mutable-ok: injects stream_options into the SDK kwargs + kwargs: dict[str, Any], # mutable-ok: SDK call kwargs built by _sdk_kwargs ) -> object: if route == ROUTE_MESSAGES: return await litellm.anthropic.messages.acreate(**kwargs) if route == ROUTE_CHAT: - if kwargs.get("stream"): - stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field - kwargs["stream_options"] = { # mutable-ok: JSON field sent to litellm.acompletion - "include_usage": True, - **stream_options, + if not kwargs.get("stream"): + return await litellm.acompletion(**kwargs) + stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field + return await litellm.acompletion( + **{ + **kwargs, + "stream_options": {"include_usage": True, **stream_options}, + STREAM_USAGE_INJECTED_FIELD: "include_usage" not in stream_options, } - return await litellm.acompletion(**kwargs) + ) return await litellm.aresponses(**kwargs) async def _call_sdk(self, route: str, body: Mapping[str, object]) -> Response: diff --git a/litellm/harness/handlers/base.py b/litellm/harness/handlers/base.py index 18b7391e988..4721e9d6ff2 100644 --- a/litellm/harness/handlers/base.py +++ b/litellm/harness/handlers/base.py @@ -4,7 +4,6 @@ from __future__ import annotations from abc import ABC, abstractmethod from collections.abc import AsyncIterator -from typing import Any from litellm.harness.context import SessionContext from litellm.harness.errors import CapabilityUnsupported @@ -38,5 +37,5 @@ class BaseHarnessHandler(ABC): async def resume(self, ctx: SessionContext, native_session_id: str) -> None: """Continue the runtime's own session on the next turn.""" - async def history(self, ctx: SessionContext) -> list[dict[str, Any]]: # mutable-ok: public history() API shape + async def history(self, ctx: SessionContext) -> list[dict[str, object]]: # mutable-ok: public history() API shape raise CapabilityUnsupported(f"Harness.{self.config.harness.name} does not expose history") diff --git a/litellm/harness/handlers/cli_handler.py b/litellm/harness/handlers/cli_handler.py index 6320a877e6b..b01f449a67d 100644 --- a/litellm/harness/handlers/cli_handler.py +++ b/litellm/harness/handlers/cli_handler.py @@ -70,10 +70,10 @@ async def send_stdin(proc: Process, data: str) -> None: async def private_dir_for(sandbox: Sandbox) -> str: - tempdir = getattr(sandbox, "tempdir", None) + tempdir: Final = sandbox.tempdir if hasattr(sandbox, "tempdir") else None if tempdir is None: raise SandboxError(f"{type(sandbox).__name__} has no tempdir(); CLI harnesses need a private config dir") - path: str = await tempdir() + path: Final = await tempdir() return path diff --git a/litellm/harness/handlers/deepagents_handler.py b/litellm/harness/handlers/deepagents_handler.py index 4a36b1a08d1..f6785323292 100644 --- a/litellm/harness/handlers/deepagents_handler.py +++ b/litellm/harness/handlers/deepagents_handler.py @@ -161,7 +161,7 @@ class DeepAgentsHandler(BaseHarnessHandler): async def history( self, ctx: SessionContext - ) -> list[dict[str, Any]]: # mutable-ok: BaseHarnessHandler.history API returns OpenAI message dicts + ) -> list[dict[str, object]]: # mutable-ok: BaseHarnessHandler.history API returns OpenAI message dicts agent, deps = self._require_agent() snapshot = await agent.aget_state(self._run_config(ctx, None)) messages = (snapshot.values or MappingProxyType({})).get("messages") or () diff --git a/litellm/harness/runtime.py b/litellm/harness/runtime.py index dc24a1f0ef7..1d946521bfa 100644 --- a/litellm/harness/runtime.py +++ b/litellm/harness/runtime.py @@ -12,11 +12,10 @@ import inspect import logging import os import uuid -from collections.abc import AsyncIterator, Callable, Coroutine, Generator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Generator, Mapping, Sequence from dataclasses import dataclass, field from types import MappingProxyType from typing import ( - Any, Final, get_args, ) @@ -331,11 +330,15 @@ async def _put_end(queue: asyncio.Queue[Event | _End], end: _End) -> None: queue.put_nowait(end) +async def _await_decision(pending: Awaitable[object]) -> object: + return await pending + + async def call_approval_handler(handler: ApprovalHandler, approval: Approval) -> None: """Run on_approval (sync in a worker thread, or async) and resolve approval.""" try: if inspect.iscoroutinefunction(handler): - decision: object = await handler(approval) + decision: object = await _await_decision(handler(approval)) else: decision = await asyncio.to_thread(handler, approval) if inspect.isawaitable(decision): @@ -749,7 +752,7 @@ class AsyncSession: async def history( self, - ) -> list[dict[str, Any]]: # mutable-ok: public API returns OpenAI-format message dicts from the handler + ) -> list[dict[str, object]]: # mutable-ok: public API returns OpenAI-format message dicts from the handler if not self.harness_config.capabilities.history: raise CapabilityUnsupported(f"Harness.{self.config.harness.name} does not expose history") await self._ensure_ready() diff --git a/litellm/images/main.py b/litellm/images/main.py index 7f94835d59b..29bac3ffe05 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -3,7 +3,7 @@ import contextvars import importlib from collections.abc import Coroutine from functools import partial -from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload +from typing import TYPE_CHECKING, Final, Literal, cast, overload from pydantic import TypeAdapter, ValidationError @@ -70,16 +70,17 @@ from litellm.utils import ( _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) # Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls -_ImageEditRequestUtils_cache: Optional["ImageEditRequestUtils"] = None +_ImageEditRequestUtils_cache: type["ImageEditRequestUtils"] | None = None -def _get_ImageEditRequestUtils() -> "ImageEditRequestUtils": +def _get_ImageEditRequestUtils() -> type["ImageEditRequestUtils"]: """Get ImageEditRequestUtils, loading it lazily if needed.""" global _ImageEditRequestUtils_cache if _ImageEditRequestUtils_cache is None: # Access via module to trigger __getattr__ if not cached - module: Final = importlib.import_module(__name__) - _ImageEditRequestUtils_cache = module.ImageEditRequestUtils + from litellm.images import main as images_main + + _ImageEditRequestUtils_cache = images_main.ImageEditRequestUtils assert _ImageEditRequestUtils_cache is not None # Type narrowing for type checker return _ImageEditRequestUtils_cache diff --git a/litellm/integrations/SlackAlerting/hanging_request_check.py b/litellm/integrations/SlackAlerting/hanging_request_check.py index 6f986144d4c..87baee91780 100644 --- a/litellm/integrations/SlackAlerting/hanging_request_check.py +++ b/litellm/integrations/SlackAlerting/hanging_request_check.py @@ -115,14 +115,16 @@ class AlertingHangingRequestCheck: if hanging_request_data.alerted: continue - request_status = await proxy_logging_obj.internal_usage_cache.async_get_cache( - key=f"request_status:{hanging_request_data.request_id}", - litellm_parent_otel_span=None, - local_only=True, - ) # this means the request status was either success or fail # and is not hanging - if request_status is not None: + if ( + await proxy_logging_obj.internal_usage_cache.async_get_cache( + key=f"request_status:{hanging_request_data.request_id}", + litellm_parent_otel_span=None, + local_only=True, + ) + is not None + ): # clear this request from hanging request cache since the request was either success or failed self.hanging_request_cache._remove_key( key=request_id, diff --git a/litellm/integrations/SlackAlerting/user_spend_alerts.py b/litellm/integrations/SlackAlerting/user_spend_alerts.py index 38794735c1b..08799a2f59f 100644 --- a/litellm/integrations/SlackAlerting/user_spend_alerts.py +++ b/litellm/integrations/SlackAlerting/user_spend_alerts.py @@ -56,8 +56,9 @@ async def fetch_user_spend_rows( today_str: Final = today.strftime("%Y-%m-%d") month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d") baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d") - raw: Final = await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) - return USER_SPEND_ROWS_ADAPTER.validate_python(raw) + return USER_SPEND_ROWS_ADAPTER.validate_python( + await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) + ) def _daily_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None: diff --git a/litellm/integrations/agentops/agentops.py b/litellm/integrations/agentops/agentops.py index 9043ab1853a..3364520a8ae 100644 --- a/litellm/integrations/agentops/agentops.py +++ b/litellm/integrations/agentops/agentops.py @@ -4,11 +4,16 @@ AgentOps integration for LiteLLM - Provides OpenTelemetry tracing for LLM calls import os from dataclasses import dataclass -from typing import Any, Final +from typing import Final + +from pydantic import TypeAdapter from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.custom_httpx.http_handler import get_httpx_client +_DECODED_JSON: Final = TypeAdapter(object) + @dataclass class AgentOpsConfig: @@ -58,9 +63,10 @@ class AgentOps(OpenTelemetry): project_id = None if config.api_key: try: - response: Final = self._fetch_auth_token(config.api_key, config.auth_endpoint) - jwt_token = response.get("token") - project_id = response.get("project_id") + response: Final = as_str_mapping(self._fetch_auth_token(config.api_key, config.auth_endpoint)) + if response is not None: + jwt_token = response.get("token") + project_id = response.get("project_id") except Exception: pass @@ -78,12 +84,9 @@ class AgentOps(OpenTelemetry): "telemetry.sdk.name": "agentops", } - if project_id: - resource_attrs["project.id"] = project_id + self.resource_attributes = {**resource_attrs, "project.id": project_id} if project_id else resource_attrs - self.resource_attributes = resource_attrs - - def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> dict[str, Any]: + def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> object: """ Fetch JWT authentication token from AgentOps API @@ -111,6 +114,6 @@ class AgentOps(OpenTelemetry): if response.status_code != 200: raise Exception(f"Failed to fetch auth token: {response.text}") - return response.json() + return _DECODED_JSON.validate_python(response.json()) finally: client.close() diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 24279b443aa..5a03b91d470 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -166,7 +166,7 @@ class ArizeLogger(OpenTelemetry): primary ``otel`` callback which handles proxy-level parent spans. """ - def set_attributes(self, span: Span, kwargs, response_obj: Any | None): + def set_attributes(self, span: Span, kwargs, response_obj: object): ArizeLogger.set_arize_attributes(span, kwargs, response_obj) @staticmethod diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 501f5749ea4..f72b29a6ad8 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -3,9 +3,10 @@ import os from datetime import datetime -from typing import Final +from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm import verbose_logger @@ -24,6 +25,18 @@ from litellm.utils import print_verbose API_BASE: Final = "https://api.braintrustdata.com/v1" +class _BraintrustProjectBody(TypedDict): + id: ReadOnly[str] + + +class _BraintrustProjectResponse(Protocol): + def json(self) -> _BraintrustProjectBody: ... + + +def _project_id(response: _BraintrustProjectResponse) -> str: + return response.json()["id"] + + def get_utc_datetime(): import datetime as dt @@ -80,8 +93,7 @@ class BraintrustLogger(CustomLogger): headers=self.headers, json={"name": project_name}, ) - project_dict: Final = response.json() - project_id: Final = project_dict["id"] + project_id: Final = _project_id(response) self._project_id_cache[project_name] = project_id return project_id except httpx.HTTPStatusError as e: @@ -100,8 +112,7 @@ class BraintrustLogger(CustomLogger): headers=self.headers, json={"name": project_name}, ) - project_dict: Final = response.json() - project_id: Final = project_dict["id"] + project_id: Final = _project_id(response) self._project_id_cache[project_name] = project_id return project_id except httpx.HTTPStatusError as e: @@ -112,18 +123,14 @@ class BraintrustLogger(CustomLogger): f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} ) - project_dict: Final = project.json() - - self.default_project_id = project_dict["id"] + self.default_project_id = _project_id(project) def create_sync_default_project_and_experiment(self): project: Final = self.global_braintrust_sync_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} ) - project_dict: Final = project.json() - - self.default_project_id = project_dict["id"] + self.default_project_id = _project_id(project) def log_success_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug("REACHES BRAINTRUST SUCCESS") diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py index 1e0413e8576..d0b25ea83e8 100644 --- a/litellm/integrations/cloudzero/cz_stream_api.py +++ b/litellm/integrations/cloudzero/cz_stream_api.py @@ -19,8 +19,9 @@ """Output modules for writing CBF data to various destinations.""" import zoneinfo +from collections.abc import Mapping from datetime import datetime, timezone -from typing import Any, Final +from typing import Final import httpx import polars as pl @@ -188,7 +189,7 @@ class CloudZeroStreamer: ) raise - def _prepare_batch_payload(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> dict[str, Any]: + def _prepare_batch_payload(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> Mapping[str, object]: """Prepare batch payload according to CloudZero AnyCost API format.""" # Convert batch_date to month for the API (YYYY-MM format) try: @@ -205,7 +206,7 @@ class CloudZeroStreamer: if record: data_records.append(record) - payload: Final = {"month": month_str, "operation": operation, "data": data_records} + payload: Final[Mapping[str, object]] = {"month": month_str, "operation": operation, "data": data_records} return payload diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1c12ac395c7..9e2dc902d55 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -352,7 +352,7 @@ class CustomGuardrail(CustomLogger): self, violation_message: str, request_data: dict[str, Any], - detection_info: dict[str, Any] | None = None, + detection_info: Mapping[str, object] | None = None, original_response: object = None, ) -> None: """ @@ -1368,7 +1368,10 @@ class CustomGuardrail(CustomLogger): if event_type is not None: guardrail_mode = event_type elif isinstance(self.event_hook, Mode): - guardrail_mode = GuardrailMode(**dict(self.event_hook.model_dump())) + guardrail_mode = GuardrailMode( + tags=copy.deepcopy(self.event_hook.tags), + default=copy.deepcopy(self.event_hook.default), + ) else: guardrail_mode = self.event_hook diff --git a/litellm/integrations/custom_secret_manager.py b/litellm/integrations/custom_secret_manager.py index e59842d409a..8fc5bc551c8 100644 --- a/litellm/integrations/custom_secret_manager.py +++ b/litellm/integrations/custom_secret_manager.py @@ -37,7 +37,6 @@ Usage: """ from abc import abstractmethod -from typing import Any import httpx @@ -154,7 +153,7 @@ class CustomSecretManager(BaseSecretManager): optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None, tags: dict | list | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Asynchronously write a secret to your custom secret manager. diff --git a/litellm/integrations/deepeval/types.py b/litellm/integrations/deepeval/types.py index eb4bb1f9b2a..a547863e26a 100644 --- a/litellm/integrations/deepeval/types.py +++ b/litellm/integrations/deepeval/types.py @@ -1,6 +1,6 @@ # Duplicate -> https://github.com/confident-ai/deepeval/blob/main/deepeval/tracing/api.py from enum import Enum -from typing import Any, ClassVar, Literal +from typing import ClassVar, Literal from pydantic import ConfigDict, Field @@ -55,7 +55,7 @@ class TraceApi(LiteLLMBaseModel): tool_spans: list[BaseApiSpan] = Field(alias="toolSpans") start_time: str = Field(alias="startTime") end_time: str = Field(alias="endTime") - metadata: dict[str, Any] | None = Field(None) + metadata: dict[str, object] | None = Field(None) tags: list[str] | None = Field(None) environment: str | None = Field(None) diff --git a/litellm/integrations/dotprompt/__init__.py b/litellm/integrations/dotprompt/__init__.py index 4cf6c0de2df..0a2921db6c7 100644 --- a/litellm/integrations/dotprompt/__init__.py +++ b/litellm/integrations/dotprompt/__init__.py @@ -61,7 +61,9 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom prompt_file: Final = getattr(litellm_params, "prompt_file", None) # Handle dotprompt_content from database - dotprompt_content: Final = getattr(litellm_params, "dotprompt_content", None) + dotprompt_content: Final = ( + litellm_params.dotprompt_content if hasattr(litellm_params, "dotprompt_content") else None + ) if dotprompt_content and not prompt_data and not prompt_file: prompt_data = get_prompt_data_from_dotprompt_content(dotprompt_content) diff --git a/litellm/integrations/dotprompt/dotprompt_manager.py b/litellm/integrations/dotprompt/dotprompt_manager.py index c646dbf4e2e..77d70dd761c 100644 --- a/litellm/integrations/dotprompt/dotprompt_manager.py +++ b/litellm/integrations/dotprompt/dotprompt_manager.py @@ -7,6 +7,8 @@ import json from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final +from typing_extensions import ReadOnly, TypedDict + from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.integrations.prompt_management_base import PromptManagementClient from litellm.types.llms.openai import AllMessageValues @@ -18,7 +20,12 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -from .prompt_manager import PromptManager, PromptTemplate +from .prompt_manager import PromptJsonEntry, PromptManager, PromptTemplate + + +class _PromptJsonData(TypedDict, total=False): + content: ReadOnly[str] + metadata: ReadOnly[dict[str, object]] class DotpromptManager(CustomPromptManagement): @@ -342,17 +349,17 @@ class DotpromptManager(CustomPromptManagement): if self._prompt_manager: self._prompt_manager.reload_prompts() - def add_prompt_from_json(self, prompt_id: str, json_data: dict[str, Any]) -> None: + def add_prompt_from_json(self, prompt_id: str, json_data: _PromptJsonData) -> None: """Add a prompt from JSON data.""" content: Final = json_data.get("content", "") metadata: Final = json_data.get("metadata", {}) self.prompt_manager.add_prompt(prompt_id, content, metadata) - def load_prompts_from_json(self, prompts_data: dict[str, dict[str, object]]) -> None: + def load_prompts_from_json(self, prompts_data: Mapping[str, object]) -> None: """Load multiple prompts from JSON data.""" self.prompt_manager.load_prompts_from_json_data(prompts_data) - def get_prompts_as_json(self) -> dict[str, dict[str, object]]: + def get_prompts_as_json(self) -> dict[str, PromptJsonEntry]: """Get all prompts in JSON format.""" return self.prompt_manager.get_all_prompts_as_json() diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index 0eda266740c..0e52ca0b62f 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -12,6 +12,8 @@ from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm.litellm_core_utils.core_helpers import as_str_mapping + class _PromptFileJson(TypedDict): """JSON form of a .prompt file: rendered template text plus its frontmatter.""" @@ -20,6 +22,11 @@ class _PromptFileJson(TypedDict): metadata: ReadOnly[NotRequired[dict[str, object]]] +class PromptJsonEntry(TypedDict): + content: ReadOnly[str] + metadata: ReadOnly[dict[str, object]] + + def strip_version_suffix(prompt_id: str) -> str | None: base, separator, version = prompt_id.rpartition(".v") if separator and base and version.isdigit(): @@ -358,20 +365,26 @@ class PromptManager: return f"---\n{frontmatter_yaml}---\n{content}" - def get_all_prompts_as_json(self) -> dict[str, dict[str, Any]]: + def get_all_prompts_as_json(self) -> dict[str, PromptJsonEntry]: """Get all loaded prompts in JSON format. Returns: Dictionary mapping prompt_id to prompt data """ - result: Final = {} - for prompt_id, template in self.prompts.items(): - result[prompt_id] = { + return { + prompt_id: { "content": template.content, "metadata": template.metadata, } - return result + for prompt_id, template in self.prompts.items() + } - def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, object]]) -> None: - """Load additional prompts from JSON data (merges with existing prompts).""" - self._load_prompts_from_json(prompt_data) + def load_prompts_from_json_data(self, prompt_data: Mapping[str, object]) -> None: + """Load additional prompts from JSON data (merges with existing prompts). Entries that are not + mappings are skipped, the same as a malformed entry inside the loader.""" + entries: Final = { + prompt_id: dict(entry) + for prompt_id, raw in prompt_data.items() + if (entry := as_str_mapping(raw)) is not None + } + self._load_prompts_from_json(entries) diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index 0ecf04326b3..cf871013265 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections.abc import Mapping from typing import Any, Final from .base import FocusDestination @@ -20,7 +21,7 @@ class FocusDestinationFactory: *, provider: str, prefix: str, - config: dict[str, Any] | None = None, + config: Mapping[str, object] | None = None, ) -> FocusDestination: """Return a destination implementation for the requested provider.""" provider_lower: Final = provider.lower() @@ -39,7 +40,7 @@ class FocusDestinationFactory: def _resolve_config( *, provider: str, - overrides: dict[str, object], + overrides: Mapping[str, object], ) -> dict[str, Any]: if provider == "s3": resolved = { diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index 615aba72954..5c345630984 100644 --- a/litellm/integrations/focus/export_engine.py +++ b/litellm/integrations/focus/export_engine.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Callable -from typing import Any, Final, SupportsFloat, SupportsIndex, SupportsInt, cast +from collections.abc import Callable, Mapping +from typing import Final, SupportsFloat, SupportsIndex, SupportsInt, cast import polars as pl @@ -24,7 +24,7 @@ class FocusExportEngine: provider: str, export_format: str, prefix: str, - destination_config: dict[str, Any] | None = None, + destination_config: Mapping[str, object] | None = None, ) -> None: self.provider = provider self.export_format = export_format @@ -77,7 +77,7 @@ class FocusExportEngine: return FocusParquetSerializer() raise NotImplementedError(f"Export format '{self.export_format}' not supported. Use 'parquet' or 'csv'.") - async def dry_run_export_usage_data(self, limit: int | None) -> dict[str, Any]: + async def dry_run_export_usage_data(self, limit: int | None) -> dict[str, object]: data: Final = await self.database.get_usage_data(limit=limit) normalized: Final = self.transformer.transform(data) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index e961b53a240..71b61296cfc 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -6,6 +6,8 @@ import os from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, ClassVar, Final, cast +from typing_extensions import TypedDict, Unpack + import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger @@ -25,6 +27,10 @@ FOCUS_USAGE_DATA_JOB_NAME: Final = "focus_export_usage_data" DEFAULT_DRY_RUN_LIMIT: Final = 500 +class _CustomLoggerOptions(TypedDict, total=False, extra_items=object): + pass + + class FocusLogger(CustomLogger): """Coordinates Focus export jobs across transformer/serializer/destination layers.""" @@ -40,7 +46,7 @@ class FocusLogger(CustomLogger): interval_seconds: int | None = None, prefix: str | None = None, destination_config: dict[str, Any] | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomLoggerOptions], ) -> None: super().__init__(**kwargs) self.provider = (provider or os.getenv("FOCUS_PROVIDER") or "s3").lower() diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 5bc7247e02a..99cd8b31b63 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -3,6 +3,7 @@ import hashlib import json import os import time +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final @@ -15,6 +16,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase from litellm.litellm_core_utils.cloud_storage_security import ( sanitize_cloud_object_component, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload @@ -134,7 +136,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): """ return f"{date_str}/batch-{batch_id}.ndjson" - def _get_config_key(self, kwargs: dict[str, Any]) -> str: + def _get_config_key(self, kwargs: Mapping[str, object]) -> str: """ Extract a synchronous grouping key from kwargs to group items by GCS config. This allows us to batch items with the same bucket/credentials together. @@ -143,7 +145,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): This key may contain sensitive information (bucket names, paths) - use _sanitize_config_key() for logging purposes. """ - standard_callback_dynamic_params: Final = kwargs.get("standard_callback_dynamic_params", None) or {} + standard_callback_dynamic_params: Final = ( + as_str_mapping(kwargs.get("standard_callback_dynamic_params", None)) or {} + ) bucket_name = standard_callback_dynamic_params.get("gcs_bucket_name", None) or self.BUCKET_NAME or "default" path_service_account: Final = ( diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py index 9dfd75e5559..3f8cf93e11c 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py @@ -9,6 +9,7 @@ Usage: """ import asyncio +import os from collections.abc import AsyncIterable, Iterable from typing import Final @@ -40,7 +41,7 @@ _mocks_initialized = False # Default mock latency in seconds (simulates network round-trip) # Typical GCS API calls take 100-300ms for uploads, 50-150ms for GET/DELETE -_MOCK_LATENCY_SECONDS: Final = float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0 +_MOCK_LATENCY_SECONDS: Final = float(os.getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0 async def _mock_async_handler_get(self, url, params=None, headers=None, follow_redirects=None): diff --git a/litellm/integrations/gcs_pubsub/pub_sub.py b/litellm/integrations/gcs_pubsub/pub_sub.py index c799245c215..cbf6898d978 100644 --- a/litellm/integrations/gcs_pubsub/pub_sub.py +++ b/litellm/integrations/gcs_pubsub/pub_sub.py @@ -12,6 +12,8 @@ import os import traceback from typing import TYPE_CHECKING, Any, Final +from pydantic import TypeAdapter + from litellm.types.utils import StandardLoggingPayload if TYPE_CHECKING: @@ -27,6 +29,8 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) +_DECODED_JSON: Final = TypeAdapter(object) + class GcsPubSubLogger(CustomBatchLogger): def __init__( @@ -152,7 +156,7 @@ class GcsPubSubLogger(CustomBatchLogger): finally: self.log_queue.clear() - async def publish_message(self, message: SpendLogsPayload | StandardLoggingPayload) -> dict[str, Any] | None: + async def publish_message(self, message: SpendLogsPayload | StandardLoggingPayload) -> object: """ Publish message to Google Cloud Pub/Sub using REST API @@ -188,7 +192,7 @@ class GcsPubSubLogger(CustomBatchLogger): raise Exception(f"Failed to publish message: {response.text}") verbose_logger.debug("Pub/Sub response: %s", response.text) - return response.json() + return _DECODED_JSON.validate_python(response.json()) except Exception as e: verbose_logger.error("Pub/Sub publish error: %s", str(e)) diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 86f80047e52..b8c6e91d0d7 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -20,6 +20,7 @@ from litellm.integrations.langfuse.langfuse_mock_client import ( should_use_langfuse_mock, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, filter_exceptions_from_params, reconstruct_model_name, safe_deep_copy, @@ -1268,11 +1269,11 @@ def log_provider_specific_information_as_span( *, tracing: "LangfuseTracing", parent: "LangfuseObservation", - enrichments: Mapping[str, Any], + enrichments: Mapping[str, object], ): """Logs provider-specific information as spans under the generation.""" - _hidden_params: Final[Mapping[str, object] | None] = enrichments.get("hidden_params", None) + _hidden_params: Final = as_str_mapping(enrichments.get("hidden_params", None)) if _hidden_params is None: return diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index f195cd97e3c..9364a3d4216 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -12,6 +12,7 @@ from litellm.integrations.langfuse.langfuse_otel_attributes import ( serialize_langfuse_datetime, ) from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.types.integrations.langfuse_otel import ( LangfuseSpanAttributes, @@ -222,7 +223,7 @@ class LangfuseOtelLogger(OpenTelemetry): ) @staticmethod - def _set_langfuse_specific_attributes(span: Span, kwargs: Mapping[str, Any], response_obj): + def _set_langfuse_specific_attributes(span: Span, kwargs: Mapping[str, object], response_obj): """ Sets Langfuse specific metadata attributes onto the OTEL span. @@ -234,7 +235,7 @@ class LangfuseOtelLogger(OpenTelemetry): from litellm.integrations.arize._utils import safe_set_attribute from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - dynamic_params: Final = kwargs.get("standard_callback_dynamic_params") + dynamic_params: Final = as_str_mapping(kwargs.get("standard_callback_dynamic_params")) langfuse_environment: Final = ( dynamic_params.get("langfuse_environment") if dynamic_params else None ) or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT") diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 3786087ba91..50ff7002068 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -2,6 +2,7 @@ Call Hook for LiteLLM Proxy which allows Langfuse prompt management. """ +from collections.abc import Mapping from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast @@ -110,6 +111,10 @@ def _remember_trace_id(litellm_call_id: object, logged: LangfuseLoggedEvent) -> in_memory_trace_id_cache.set_cache(litellm_call_id=litellm_call_id, service_name="langfuse", trace_id=trace_id) +def _params_without_model(config: Mapping[str, object]) -> Mapping[str, object]: + return {k: v for k, v in config.items() if k != "model"} + + class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogger): def __init__( self, @@ -182,12 +187,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge return compiled_prompt def _get_optional_params_from_langfuse(self, langfuse_prompt_client: PROMPT_CLIENT) -> dict: - config: Final = langfuse_prompt_client.config - optional_params: Final = {} - for k, v in config.items(): - if k != "model": - optional_params[k] = v - return optional_params + return dict(_params_without_model(langfuse_prompt_client.config)) async def async_get_chat_completion_prompt( self, diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 01812149498..ff1f3568685 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -7,10 +7,11 @@ import random import traceback import types from datetime import datetime, timezone -from typing import Any, Final +from typing import Final import httpx from pydantic import BaseModel, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -31,6 +32,10 @@ from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPa _RUN_RESPONSE: Final = TypeAdapter(object) +class _DynamicParamsKwargs(TypedDict, total=False): + standard_callback_dynamic_params: ReadOnly[StandardCallbackDynamicParams | None] + + def is_serializable(value): non_serializable_types: Final = ( types.CoroutineType, @@ -474,7 +479,7 @@ class LangsmithLogger(CustomBatchLogger): return log_queue_by_credentials - def _get_sampling_rate_to_use_for_request(self, kwargs: dict[str, Any]) -> float: + def _get_sampling_rate_to_use_for_request(self, kwargs: _DynamicParamsKwargs) -> float: standard_callback_dynamic_params: Final[StandardCallbackDynamicParams | None] = kwargs.get( "standard_callback_dynamic_params", None ) @@ -485,7 +490,7 @@ class LangsmithLogger(CustomBatchLogger): sampling_rate = float(_sampling_rate) return sampling_rate - def _get_credentials_to_use_for_request(self, kwargs: dict[str, Any]) -> LangsmithCredentialsObject: + def _get_credentials_to_use_for_request(self, kwargs: _DynamicParamsKwargs) -> LangsmithCredentialsObject: """ Handles key/team based logging diff --git a/litellm/integrations/langtrace.py b/litellm/integrations/langtrace.py index eeff76e7063..ea411baee7e 100644 --- a/litellm/integrations/langtrace.py +++ b/litellm/integrations/langtrace.py @@ -6,7 +6,7 @@ from litellm.types.integrations.otel_span_attributes import SpanAttributes if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any diff --git a/litellm/integrations/levo/levo.py b/litellm/integrations/levo/levo.py index 4d2b4edf3cd..c1852e9bfe4 100644 --- a/litellm/integrations/levo/levo.py +++ b/litellm/integrations/levo/levo.py @@ -4,8 +4,6 @@ from typing import TYPE_CHECKING, Any, Final from litellm.integrations.opentelemetry import OpenTelemetry if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - from litellm.integrations.opentelemetry import ( OpenTelemetryConfig as _OpenTelemetryConfig, ) @@ -13,11 +11,9 @@ if TYPE_CHECKING: Protocol = _Protocol OpenTelemetryConfig = _OpenTelemetryConfig - Span = _Span | Any else: Protocol = Any OpenTelemetryConfig = Any - Span = Any class LevoConfig: diff --git a/litellm/integrations/logfire_logger.py b/litellm/integrations/logfire_logger.py index 848369a1cbb..862893fa883 100644 --- a/litellm/integrations/logfire_logger.py +++ b/litellm/integrations/logfire_logger.py @@ -3,8 +3,9 @@ import os import traceback +from collections.abc import Mapping from enum import Enum -from typing import Any, Final, NamedTuple +from typing import Final, NamedTuple from typing_extensions import LiteralString @@ -15,7 +16,7 @@ from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info class SpanConfig(NamedTuple): message_template: LiteralString - span_data: dict[str, Any] + span_data: Mapping[str, object] class LogfireLevel(str, Enum): diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py index 9829ef4e18f..3623335f599 100644 --- a/litellm/integrations/newrelic/newrelic.py +++ b/litellm/integrations/newrelic/newrelic.py @@ -49,11 +49,12 @@ import time import uuid from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Any, Final +from typing import Final, Protocol import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.newrelic import NewRelicInitParams @@ -65,6 +66,10 @@ except ImportError: _newrelic_agent = None +class _SupportsTimestamp(Protocol): + def timestamp(self) -> float: ... + + class NewRelicLogger(CustomLogger): """ New Relic logger for LiteLLM to send AI monitoring events. @@ -126,13 +131,13 @@ class NewRelicLogger(CustomLogger): verbose_logger.error("Failed to initialize New Relic agent: %s. Integration will be disabled.", e) self.enabled = False - def _get_newrelic_params(self) -> dict: + def _get_newrelic_params(self) -> Mapping[str, object]: """ Get the newrelic_params from litellm.newrelic_params These are params specific to initializing the NewRelicLogger e.g. turn_off_message_logging """ - dict_newrelic_params: dict = {} + dict_newrelic_params: Mapping[str, object] = {} if litellm.newrelic_params is not None: if isinstance(litellm.newrelic_params, NewRelicInitParams): dict_newrelic_params = litellm.newrelic_params.model_dump() @@ -401,9 +406,9 @@ class NewRelicLogger(CustomLogger): return choices[0].get("finish_reason") or "unknown" return "unknown" - def _to_epoch_ms(self, t: Any) -> float: + def _to_epoch_ms(self, t: datetime | float) -> float: """Convert a datetime or float timestamp to epoch milliseconds.""" - if hasattr(t, "timestamp"): + if has_attributes(t, _SupportsTimestamp, "timestamp"): return t.timestamp() * 1000.0 return float(t) * 1000.0 @@ -518,13 +523,13 @@ class NewRelicLogger(CustomLogger): sequence = 0 # Extract timestamps, preferring StandardLoggingPayload - start_time = None + start_time: datetime | float | None = None if standard_logging_object: start_time = standard_logging_object.get("startTime") if not start_time: start_time = kwargs.get("start_time") - end_time = None + end_time: datetime | float | None = None if standard_logging_object: end_time = standard_logging_object.get("endTime") if not end_time: diff --git a/litellm/integrations/newrelic/newrelic_metrics.py b/litellm/integrations/newrelic/newrelic_metrics.py index a2cedeba0f5..641d1d86188 100644 --- a/litellm/integrations/newrelic/newrelic_metrics.py +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -254,7 +254,7 @@ class NewRelicMetricsLogger(CustomBatchLogger): if not self.log_queue: return if _pass < NEWRELIC_METRICS_MAX_DRAIN_PASSES - 1: - await asyncio.sleep(2**_pass) + await asyncio.sleep(1 << _pass) async with self.flush_lock: tried_ids: Final = frozenset(id(record) for record in attempted) survivors: Final = tuple(record for record in self.log_queue if id(record) not in tried_ids) diff --git a/litellm/integrations/opik/opik_payload_builder/payload_builders.py b/litellm/integrations/opik/opik_payload_builder/payload_builders.py index 3aaf5bfc162..2909ab9d63c 100644 --- a/litellm/integrations/opik/opik_payload_builder/payload_builders.py +++ b/litellm/integrations/opik/opik_payload_builder/payload_builders.py @@ -57,9 +57,7 @@ def build_span_payload( span_id: Final = utils.create_uuid7() model: Final[str] = response_obj.get("model", "unknown-model") - obj_type: Final[str] = response_obj.get("object", "unknown-object") - created: Final[int] = response_obj.get("created", 0) - span_name: Final = f"{model}_{obj_type}_{created}" + span_name: Final = f"{model}_{response_obj.get('object', 'unknown-object')}_{response_obj.get('created', 0)}" _logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id) diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index cca12b885c1..dd150082252 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -130,7 +130,8 @@ class _EnvWithoutBareExcludedServices(PydanticBaseSettingsSource): return self._env_settings.get_field_value(field, field_name) def __call__(self) -> dict[str, object]: - return {key: value for key, value in self._env_settings().items() if key != "excluded_services"} + settings: Final = self._env_settings() + return {key: settings[key] for key in settings if key != "excluded_services"} class OpenTelemetryV2Config(BaseSettings): diff --git a/litellm/integrations/otel/model/destination.py b/litellm/integrations/otel/model/destination.py index 4226b4361dd..e760ba3600b 100644 --- a/litellm/integrations/otel/model/destination.py +++ b/litellm/integrations/otel/model/destination.py @@ -42,7 +42,7 @@ class OtelDestination(LiteLLMBaseModel): success_sampling_rate: float | None = Field( default=None, description=( - "Share of the request trees forwarded, 0.0..1.0, drawn once per request when its root span ends; " + "Share of the request trees forwarded, 0.0..1.0, picked by a hash of the trace id; " "``None`` forwards every one." ), ) diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index 5d37894bb8f..a69c334a1fb 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -332,7 +332,7 @@ def caller_session_id(kwargs: Mapping[str, object], trace: TraceControls) -> str return (None if echoes_trace_id else explicit) or trace.session_id or None -def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None: +def time_to_first_chunk_seconds(kwargs: Mapping[str, object]) -> float | None: """Seconds from the upstream request being issued (``api_call_start_time``) to the first streamed chunk (``completion_start_time``); ``None`` for non-streaming calls, where ``completion_start_time`` is backfilled with the @@ -340,8 +340,8 @@ def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None: optional_params: Final = cast(Mapping[str, object], kwargs.get("optional_params") or {}) if not optional_params.get("stream"): return None - api_call_start: Final = to_seconds(kwargs.get("api_call_start_time")) - completion_start: Final = to_seconds(kwargs.get("completion_start_time")) + api_call_start: Final = _epoch_seconds(kwargs.get("api_call_start_time")) + completion_start: Final = _epoch_seconds(kwargs.get("completion_start_time")) if api_call_start is None or completion_start is None: return None return completion_start - api_call_start diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index d11505ce026..9c62c45e489 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -7,7 +7,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType -from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, cast from urllib.parse import urlsplit from typing_extensions import ReadOnly, TypedDict @@ -716,6 +716,15 @@ def _json_or_none(value: object) -> str | None: return None +class _EnumValue(Protocol): + @property + def value(self) -> object: ... + + +def _enum_value(item: _EnumValue) -> object: + return item.value + + def _guardrail_mode_str(value: object) -> str | None: """Normalize ``guardrail_mode`` to a stable string. @@ -732,12 +741,12 @@ def _guardrail_mode_str(value: object) -> str | None: for item in value: if item is None: continue - part = as_str(item.value) if isinstance(item, Enum) else as_str(item) + part = as_str(_enum_value(item)) if isinstance(item, Enum) else as_str(item) if part: parts.append(part) return ",".join(parts) or None if isinstance(value, Enum): - return as_str(value.value) + return as_str(_enum_value(value)) return as_str(value) diff --git a/litellm/integrations/otel/mount.py b/litellm/integrations/otel/mount.py index cb52ad24c8e..c0d42eeb5c5 100644 --- a/litellm/integrations/otel/mount.py +++ b/litellm/integrations/otel/mount.py @@ -107,8 +107,8 @@ def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | No Safe no-op when the V2 gate is off or ``opentelemetry-instrumentation-fastapi`` is unavailable. This MUST be called at app-creation time — once the lifespan - runs, the middleware stack is frozen and ``instrument_app`` raises "Cannot add - middleware after an application has started". + runs, the middleware stack is frozen and middleware added later never runs, so + a late call is skipped without registering anything. No ``TracerProvider`` is passed, so the instrumentation binds to the OTel global ``ProxyTracerProvider``; the proxy publishes the real provider as the global @@ -126,6 +126,9 @@ def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | No try: if not is_otel_v2_enabled(): return + if app.middleware_stack is not None: + verbose_logger.debug("Skipping OTel V2 FastAPI instrumentation: the app has already started") + return # Lazy: only the V2-enabled path needs the optional # ``opentelemetry-instrumentation-fastapi`` package, which is not part of the diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 7ec3711ee65..72d51471ba2 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -173,6 +173,7 @@ class _MetricKwargs(TypedDict, total=False): response_cost: ReadOnly[float | None] completion_start_time: ReadOnly[datetime | float | str | None] api_call_start_time: ReadOnly[datetime | float | str | None] + end_time: ReadOnly[datetime | float | str | None] def resolve_error_type(kwargs: Mapping[str, Any]) -> str: @@ -414,15 +415,12 @@ class GenAIMetricRecorder: if generation_time > 0: self._metrics.time_per_output_token.record(generation_time / completion_tokens, attributes=common_attrs) - def _record_response_duration(self, kwargs: Mapping[str, Any], end_time: datetime, common_attrs: dict) -> None: + def _record_response_duration(self, kwargs: _MetricKwargs, end_time: datetime, common_attrs: dict) -> None: api_call_start_time: Final = kwargs.get("api_call_start_time") if api_call_start_time is None: return - _end_time = kwargs.get("end_time") or end_time - if _end_time is None: - _end_time = datetime.now() api_call_start: Final = to_seconds(api_call_start_time) - end_ts: Final = to_seconds(_end_time) + end_ts: Final = to_seconds(kwargs.get("end_time") or end_time) if api_call_start is None or end_ts is None: return duration: Final = end_ts - api_call_start diff --git a/litellm/integrations/otel/plumbing/otlp_json.py b/litellm/integrations/otel/plumbing/otlp_json.py index bc6d7d435b0..fc0351d5049 100644 --- a/litellm/integrations/otel/plumbing/otlp_json.py +++ b/litellm/integrations/otel/plumbing/otlp_json.py @@ -6,6 +6,7 @@ retry loop and swaps the payload for OTLP/JSON (enums as integers, ids as hex). import base64 import json +import logging from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final, TypeAlias @@ -13,9 +14,13 @@ from typing import Final, TypeAlias import requests from google.protobuf.json_format import MessageToDict from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.exporter.otlp.proto.http._common import RequestPayloadTooLargeError from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SpanExportResult +from opentelemetry.semconv.attributes.http_attributes import HTTP_RESPONSE_STATUS_CODE +_logger: Final = logging.getLogger(__name__) JSON_CONTENT_TYPE: Final = "application/json" _HEX_ID_KEYS: Final = frozenset({"traceId", "spanId", "parentSpanId"}) @@ -66,12 +71,49 @@ class OTLPJsonSpanExporter(OTLPSpanExporter): def __init__( self, endpoint: str | None, - headers: dict[str, str], # mutable-ok: SDK __init__ takes Dict + headers: Mapping[str, str], certificate_file: str | None = None, session: "requests.Session | None" = None, + *, + max_request_size: int | None = None, ) -> None: - super().__init__(endpoint=endpoint, headers=headers, certificate_file=certificate_file, session=session) - self._session.headers["Content-Type"] = JSON_CONTENT_TYPE + super().__init__( + endpoint=endpoint, + headers={**{key.lower(): value for key, value in headers.items()}, "content-type": JSON_CONTENT_TYPE}, + certificate_file=certificate_file, + session=session, + max_request_size=max_request_size, + ) - def _serialize_spans(self, spans: Sequence[ReadableSpan]) -> bytes: - return encode_spans_json(spans) + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + if self._shutdown: + _logger.warning("Exporter already shutdown, ignoring batch") + return SpanExportResult.FAILURE + + with self._metrics.export_operation(len(spans)) as result: + try: + serialized_data: Final = encode_spans_json(spans) + except Exception as error: # noqa: BLE001 # SDK export contract reports encoding failures as results + _logger.error("Failed to encode span batch: %s", error) + result.error = error + return SpanExportResult.FAILURE + + if 0 < self._max_request_size < len(serialized_data): + result.error = RequestPayloadTooLargeError( + f"Serialized span request size {len(serialized_data)} bytes " + f"exceeds max_request_size {self._max_request_size} bytes." + ) + _logger.warning("Dropping span batch: %s", result.error) + return SpanExportResult.FAILURE + + export_result: Final = self._client.export(serialized_data) + if not export_result.success: + result.error = export_result.error + result.error_attrs = ( + {HTTP_RESPONSE_STATUS_CODE: export_result.status_code} + if export_result.status_code is not None + else None + ) + return SpanExportResult.FAILURE + + return SpanExportResult.SUCCESS diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 898eab5715b..fc1e9b7099d 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -1,13 +1,14 @@ """Provider / exporter factory + the Baggage span processor.""" import queue -import random import threading import time +import weakref from collections import OrderedDict from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from hashlib import sha256 from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from opentelemetry import _logs, baggage, metrics, trace from opentelemetry._logs import Logger, LoggerProvider, NoOpLoggerProvider @@ -486,15 +487,15 @@ def is_llm_call_span(span: ReadableSpan) -> bool: return GenAI.OPERATION_NAME in attributes and MCP.METHOD_NAME not in attributes -#: Request trees held back for a sampled destination until the last span of the trace +#: Request trees held back for a sampled destination until the last span of the request #: still open ends, the verdicts kept for the spans that start after that (a post-call -#: database write), and the spans in flight per trace. All bounded, so a span that never -#: ends, or a flood of requests, cannot hold spans for ever: the oldest tree is decided -#: on what it has. +#: database write), the spans in flight per request, and the request each open span +#: belongs to. All bounded, so a span that never ends, or a flood of requests, cannot +#: hold spans for ever: the oldest tree is decided on what it has. _MAX_PENDING_TREES: Final = 1024 _MAX_PENDING_SPANS_PER_TREE: Final = 512 _MAX_REMEMBERED_VERDICTS: Final = 4096 -_MAX_OPEN_TRACES: Final = 16384 +_MAX_OPEN_REQUESTS: Final = 16384 class _PendingTree: @@ -512,25 +513,35 @@ def _is_sampled(destination: "OtelDestination") -> bool: return destination.success_sampling_rate is not None or destination.error_sampling_rate is not None -#: A verdict is one request tree's draw for one destination at its rates: the exporter +#: A request is its trace id and the id of its root, the span of it whose parent is +#: remote or missing. Calls sharing a caller's trace are separate requests, each decided +#: on its own outcome. +_RequestKey: TypeAlias = tuple[int, int] + +#: A verdict is one request's keep or drop for one destination at its rates: the exporter #: ``cache_key`` leaves the rates out, so two destinations sharing an exporter with -#: different rates would otherwise share one draw. -_VerdictKey = tuple[int, object, float | None, float | None] +#: different rates would otherwise share one verdict. +_VerdictKey: TypeAlias = tuple[int, int, object, float | None, float | None] -def _verdict_key(trace_id: int, destination: "OtelDestination") -> _VerdictKey: - return (trace_id, destination.cache_key(), destination.success_sampling_rate, destination.error_sampling_rate) +def _verdict_key(request: _RequestKey, destination: "OtelDestination") -> _VerdictKey: + return (*request, destination.cache_key(), destination.success_sampling_rate, destination.error_sampling_rate) -def _keeps_tree(destination: "OtelDestination", failed: bool, draw: Callable[[], float]) -> bool: - """Whether one draw against the destination's rate keeps a request tree, read the - way the legacy Arize callback reads ``arize_success_sampling_rate`` / - ``arize_error_sampling_rate``: a tree with a failed span answers to the error - rate, an unset rate keeps everything and ``0.0`` keeps nothing.""" +def _keeps_tree(destination: "OtelDestination", failed: bool, trace_id: int) -> bool: + """Whether the destination's rate keeps a request tree of ``trace_id``: a tree with + a failed span answers to ``arize_error_sampling_rate``, the rest to + ``arize_success_sampling_rate``, an unset rate keeps everything and ``0.0`` keeps + nothing. + + The first eight bytes of a SHA-256 of the trace id are read against the rate, so + every process deciding a request of the trace with the same outcome reaches the same + verdict.""" rate: Final = destination.error_sampling_rate if failed else destination.success_sampling_rate if rate is None: return True - return rate > 0.0 and draw() <= rate + digest: Final = sha256(trace_id.to_bytes(16, "big")).digest() + return int.from_bytes(digest[:8], "big") < round(rate * 2**64) def _in_scope(span: ReadableSpan, scope: "OtelSpanScope") -> bool: @@ -667,10 +678,8 @@ class TenantFanOutSpanProcessor(SpanProcessor): default_capture_by_backend: Mapping[str, str] = MappingProxyType({}), pending_drains: int = _MAX_PENDING_DRAINS, drain_pool: _DrainPool | None = None, - sampling_draw: Callable[[], float] | None = None, ) -> None: self._operator_sinks: Final = operator_sinks - self._draw: Final = sampling_draw if sampling_draw is not None else random.random self._excluded_db_systems: Final = excluded_db_systems self._default_capture: Final = default_capture self._default_capture_by_backend: Final = default_capture_by_backend @@ -682,24 +691,32 @@ class TenantFanOutSpanProcessor(SpanProcessor): self._processors: OrderedDict[object, SpanProcessor] = OrderedDict() # mutable-ok: bounded LRU self._retired: OrderedDict[int, SpanProcessor] = OrderedDict() # mutable-ok: drains as exports finish self._exporting: dict[int, int] = {} # mutable-ok: per-processor in-flight export count - self._pending: OrderedDict[int, _PendingTree] = OrderedDict() # mutable-ok: bounded, by trace id - self._open: OrderedDict[int, int] = OrderedDict() # mutable-ok: bounded, spans in flight by trace id + self._pending: OrderedDict[_RequestKey, _PendingTree] = OrderedDict() # mutable-ok: bounded, by request + self._open: OrderedDict[_RequestKey, set[int]] = OrderedDict() # mutable-ok: bounded, open span ids by request + self._roots: dict[int, _RequestKey] = {} # mutable-ok: holds only the spans ``_open`` holds, request by span id + self._span_roots: weakref.WeakKeyDictionary[SDKSpan, int] = weakref.WeakKeyDictionary() self._verdicts: OrderedDict[_VerdictKey, bool] = OrderedDict() # mutable-ok: bounded LRU self._drain: Final = drain_pool if drain_pool is not None else _DrainPool(capacity=pending_drains) def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None: context: Final = span.context - trace_id: Final = context.trace_id if context is not None else 0 with self._lock: - self._open[trace_id] = self._open.get(trace_id, 0) + 1 - while len(self._open) > _MAX_OPEN_TRACES: - self._open.popitem(last=False) + request: Final = self._request_of(span, trace.get_current_span(parent_context)) + if context is not None: + self._span_roots[span] = request[1] + self._roots[context.span_id] = request + self._open.setdefault(request, set()).add(context.span_id) + while len(self._open) > _MAX_OPEN_REQUESTS: + for forgotten in self._open.popitem(last=False)[1]: + self._roots.pop(forgotten, None) def on_end(self, span: ReadableSpan) -> None: suppressed: Final = suppressed_backends() attributes: Final = span.attributes or _NO_ATTRIBUTES - context: Final = span.context - trace_id: Final = context.trace_id if context is not None else 0 + with self._lock: + request: Final = self._request_of(span) + if span.context is not None: + self._roots.pop(span.context.span_id, None) failed: Final = span.status.status_code is StatusCode.ERROR for destination in request_destinations(): if ( @@ -716,11 +733,32 @@ class TenantFanOutSpanProcessor(SpanProcessor): if not _is_sampled(destination): self._forward(span, destination) continue - for held_span, held_destination in self._route(trace_id, span, destination, failed): + for held_span, held_destination in self._route(request, span, destination, failed): self._forward(held_span, held_destination) - for held_span, held_destination in self._settle(trace_id, failed): + for held_span, held_destination in self._settle(request, span, failed): self._forward(held_span, held_destination) + def _request_of(self, span: ReadableSpan, parent_span: Span | None = None) -> _RequestKey: + """The request ``span`` belongs to, read under ``_lock``: the request recorded while + it is open, else the span itself when its parent is remote or missing, else its + parent's request, recorded while the parent is open or on ``parent_span`` for a + child starting under a parent that has ended, else the parent itself.""" + context: Final = span.context + if context is None: + return (0, 0) + recorded: Final = self._roots.get(context.span_id) + if recorded is not None: + return recorded + parent: Final = span.parent + if parent is None or parent.is_remote: + return (context.trace_id, context.span_id) + parent_request: Final = self._roots.get(parent.span_id) + if parent_request is not None: + return parent_request + if isinstance(parent_span, SDKSpan): + return (context.trace_id, self._span_roots.get(parent_span, parent.span_id)) + return (context.trace_id, parent.span_id) + def _default_capture_for(self, destination: "OtelDestination") -> str: backend: Final = destination.callback_name return ( @@ -743,14 +781,14 @@ class TenantFanOutSpanProcessor(SpanProcessor): self._release(processor) def _route( - self, trace_id: int, span: ReadableSpan, destination: "OtelDestination", failed: bool + self, request: _RequestKey, span: ReadableSpan, destination: "OtelDestination", failed: bool ) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": """The spans to forward now for ``span`` on a sampled ``destination``: the span itself when its tree's verdict is already known and keeps it, nothing while the - tree is held for the rest of its trace, or what the bounds made the fan-out decide + tree is held for the rest of its request, or what the bounds made the fan-out decide early (a tree that outgrew its cap is decided on what it has, and so is the oldest tree once too many are waiting, so a span that never ends holds nothing back for - ever). A failed span marks its tree before any such early draw. Once shutdown + ever). A failed span marks its tree before any such early decision. Once shutdown has started flushing, nothing is held: a span ending then is decided at once, since no later flush would reach it. @@ -758,58 +796,62 @@ class TenantFanOutSpanProcessor(SpanProcessor): another thread between the two would leave this span in a fresh tree. """ with self._lock: - verdict: Final = self._verdicts.get(_verdict_key(trace_id, destination)) + verdict: Final = self._verdicts.get(_verdict_key(request, destination)) if verdict is not None: return ((span, destination),) if verdict else () - tree: Final = self._pending.setdefault(trace_id, _PendingTree()) + tree: Final = self._pending.setdefault(request, _PendingTree()) tree.held.append((span, destination)) if failed: tree.failed = True if self._draining or len(tree.held) >= _MAX_PENDING_SPANS_PER_TREE: - return self._decide_locked(trace_id) + return self._decide_locked(request) if len(self._pending) > _MAX_PENDING_TREES: oldest: Final = next(iter(self._pending)) return self._decide_locked(oldest) return () - def _decide(self, trace_id: int) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + def _decide(self, request: _RequestKey) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": with self._lock: - return self._decide_locked(trace_id) + return self._decide_locked(request) - def _settle(self, trace_id: int, failed: bool) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": - """What the trace's held tree exports once the span that just ended was the last - of its trace still open. The model-call span ends in the post-call callback, - after the server span, so a tree is decided when its trace goes quiet rather than + def _settle( + self, request: _RequestKey, span: ReadableSpan, failed: bool + ) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + """What the request's held tree exports once the span that just ended was the last + of its request still open. The model-call span ends in the post-call callback, + after the server span, so a tree is decided when its request goes quiet rather than when its root ends, and a model call that fails late still answers to the error - rate. A trace whose starts this fan-out never saw (attached mid-flight, or forgotten - by the bound) is decided as each of its spans ends. + rate. A request whose starts this fan-out never saw (attached mid-flight, or + forgotten by the bound) is decided as each of its spans ends. """ with self._lock: - tree: Final = self._pending.get(trace_id) + tree: Final = self._pending.get(request) if tree is not None and failed: tree.failed = True - still_open: Final = self._open.pop(trace_id, 1) - 1 - if still_open > 0: - self._open[trace_id] = still_open + still_open: Final = self._open.pop(request, set()) + if span.context is not None: + still_open.discard(span.context.span_id) + if still_open: + self._open[request] = still_open return () - return self._decide_locked(trace_id) + return self._decide_locked(request) - def _decide_locked(self, trace_id: int) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": - """Draw once per destination for the held tree; the spans the draws keep, and a - verdict remembered for the spans of this trace that start later.""" - tree: Final = self._pending.pop(trace_id, None) + def _decide_locked(self, request: _RequestKey) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + """Decide once per destination for the held tree; the spans the verdicts keep, and + a verdict remembered for the spans of this request that start later.""" + tree: Final = self._pending.pop(request, None) if tree is None: return () - verdicts: dict[_VerdictKey, bool] = {} # mutable-ok: one draw per destination of this tree + verdicts: dict[_VerdictKey, bool] = {} # mutable-ok: one verdict per destination of this tree for _, destination in tree.held: - key = _verdict_key(trace_id, destination) + key = _verdict_key(request, destination) if key not in verdicts: - verdicts[key] = _keeps_tree(destination, tree.failed, self._draw) + verdicts[key] = _keeps_tree(destination, tree.failed, request[0]) self._verdicts[key] = verdicts[key] while len(self._verdicts) > _MAX_REMEMBERED_VERDICTS: self._verdicts.popitem(last=False) return tuple( - (span, destination) for span, destination in tree.held if verdicts[_verdict_key(trace_id, destination)] + (span, destination) for span, destination in tree.held if verdicts[_verdict_key(request, destination)] ) def _operator_already_writes( @@ -850,8 +892,8 @@ class TenantFanOutSpanProcessor(SpanProcessor): with self._lock: self._draining = True undecided: Final = tuple(self._pending) - for trace_id in undecided: - for held_span, held_destination in self._decide(trace_id): + for request in undecided: + for held_span, held_destination in self._decide(request): self._forward(held_span, held_destination) deadline: Final = time.monotonic() + self._drain_seconds with self._lock: diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 3befb73b0cc..141162bdfa5 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -76,6 +76,7 @@ def agentops_preset( def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" + import requests from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, ) @@ -89,7 +90,8 @@ def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: """ def __init__(self, *, endpoint: str | None, api_key: str | None) -> None: - super().__init__(endpoint=endpoint) + self._agentops_session = requests.Session() + super().__init__(endpoint=endpoint, session=self._agentops_session) self._agentops_api_key = api_key self._auth_resolved = False @@ -102,10 +104,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: try: token: Final = _fetch_agentops_jwt(self._agentops_api_key).get("token") if token: - # ``_session`` is the requests.Session the base exporter - # POSTs through; updating its Authorization header is how the - # minted JWT reaches every subsequent export. - self._session.headers["Authorization"] = f"Bearer {token}" + self._agentops_session.headers["Authorization"] = f"Bearer {token}" except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index d787375ca3c..83f85232182 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -13,6 +13,8 @@ import re import traceback from typing import Final +from pydantic import TypeAdapter + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( @@ -32,6 +34,7 @@ from litellm.types.utils import StandardLoggingPayload from .custom_batch_logger import CustomBatchLogger +_DECODED_JSON: Final = TypeAdapter(object) _BASE64_INLINE_PATTERN: Final = re.compile( r"data:(?:application|image|audio|video)/[a-zA-Z0-9.+-]+;base64,[A-Za-z0-9+/=\s]+", re.MULTILINE, @@ -270,7 +273,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): if self.sqs_queue_url is None: raise ValueError("sqs_queue_url not set") - json_data: Final = json.loads(safe_dumps(payload)) + json_data: Final = _DECODED_JSON.validate_python(json.loads(safe_dumps(payload))) if self.app_crypto: aad_bytes: Final = self.sqs_app_encryption_aad.encode("utf-8") if self.sqs_app_encryption_aad else None encrypted: Final = self.app_crypto.encrypt_json(json_data, aad=aad_bytes) diff --git a/litellm/interactions/litellm_responses_transformation/handler.py b/litellm/interactions/litellm_responses_transformation/handler.py index 8fee0fcd1b5..d84f0b5ffbd 100644 --- a/litellm/interactions/litellm_responses_transformation/handler.py +++ b/litellm/interactions/litellm_responses_transformation/handler.py @@ -116,10 +116,7 @@ class LiteLLMResponsesInteractionsHandler: """Async handler for interactions API requests.""" # Call litellm.aresponses() # Note: litellm.aresponses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] - aresponses_fn: Final[ - Callable[..., Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]] - ] = vars(litellm)["aresponses"] - responses_response: Final = await aresponses_fn( + responses_response: Final = await litellm.aresponses( **responses_request, ) diff --git a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py index dab12d447df..46e633d2b96 100644 --- a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py +++ b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py @@ -3,8 +3,8 @@ Streaming iterator for transforming Responses API stream to Interactions API str """ from collections import deque -from collections.abc import AsyncIterator, Iterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Iterator, Mapping +from typing import Final, cast from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, @@ -44,7 +44,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: request_input: InteractionInput | None, optional_params: InteractionsAPIOptionalRequestParams, custom_llm_provider: str | None = None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, ): self.model = model self.responses_stream_iterator = litellm_custom_stream_wrapper diff --git a/litellm/litellm_core_utils/app_crypto.py b/litellm/litellm_core_utils/app_crypto.py index 4ef3f34f3b8..da844620dce 100644 --- a/litellm/litellm_core_utils/app_crypto.py +++ b/litellm/litellm_core_utils/app_crypto.py @@ -4,6 +4,9 @@ import os from typing import Final from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from pydantic import TypeAdapter + +_DECODED_JSON: Final = TypeAdapter(object) class AppCrypto: @@ -12,7 +15,7 @@ class AppCrypto: raise ValueError("Master key must be 32 bytes for AES-256-GCM") self.key = master_key - def encrypt_json(self, data: dict, aad: bytes | None = None) -> dict: + def encrypt_json(self, data: object, aad: bytes | None = None) -> dict: aes: Final = AESGCM(self.key) nonce: Final = os.urandom(12) plaintext: Final = json.dumps(data).encode("utf-8") @@ -24,10 +27,10 @@ class AppCrypto: "tag": base64.b64encode(tag).decode(), } - def decrypt_json(self, enc: dict, aad: bytes | None = None) -> dict: + def decrypt_json(self, enc: dict, aad: bytes | None = None) -> object: aes: Final = AESGCM(self.key) nonce: Final = base64.b64decode(enc["nonce"]) ct: Final = base64.b64decode(enc["ciphertext"]) tag: Final = base64.b64decode(enc["tag"]) data: Final = aes.decrypt(nonce, ct + tag, aad) - return json.loads(data.decode()) + return _DECODED_JSON.validate_python(json.loads(data.decode())) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index e9689bd7d8a..8e949694881 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -3,10 +3,11 @@ import re from collections.abc import Mapping, Sequence from itertools import accumulate, repeat from types import MappingProxyType -from typing import Any, Final, cast +from typing import Final from urllib.parse import quote, unquote from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_str_mapping VERTEX_AI_MANAGED_GCS_PREFIX: Final = "litellm-vertex-files/" BEDROCK_MANAGED_S3_BATCH_PREFIX: Final = "litellm-bedrock-files-" @@ -141,7 +142,8 @@ def should_allow_legacy_cloud_file_ids( if isinstance(litellm_params, Mapping): trusted_model_credentials: Final = litellm_params.get("_litellm_internal_model_credentials") if isinstance(trusted_model_credentials, _MAPPING_PROXY_TYPE): - value = cast(Mapping[str, Any], trusted_model_credentials).get("allow_legacy_cloud_file_ids") + trusted_mapping: Final = as_str_mapping(trusted_model_credentials) + value = trusted_mapping.get("allow_legacy_cloud_file_ids") if trusted_mapping is not None else None if isinstance(value, bool): return value diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index b7c6f805233..00c6ce97e78 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -5,10 +5,11 @@ import logging import re from collections.abc import Collection, Iterable, Iterator, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable import httpx from pydantic import TypeAdapter, ValidationError +from typing_extensions import TypeIs # noqa: TID251 # narrows a duck-typed value to the Protocol its caller reads import litellm from litellm._logging import verbose_logger @@ -58,6 +59,13 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, Mapping) else None +def decoded_json_object(value: object) -> Mapping[str, object]: + mapping: Final = as_str_mapping(value) + if mapping is not None: + return mapping + raise AttributeError(f"{type(value).__name__!r} object has no attribute 'get'", name="get", obj=value) + + @runtime_checkable class _SupportsIter(Protocol): def __iter__(self) -> Iterator[object]: ... @@ -72,6 +80,24 @@ def attribute_or_none(source: object, name: str) -> object: return getattr(source, name, None) +_Shape: Final = TypeVar("_Shape") + + +def has_attributes( + value: object, + shape: type[_Shape], + *names: str, +) -> TypeIs[_Shape]: # guard-ok: hasattr sees what __getattr__ serves, which Protocol isinstance skips on 3.12+ + return all(hasattr(value, name) for name in names) + + +def require_attributes(value: object, shape: type[_Shape], *names: str) -> _Shape: + if has_attributes(value, shape, *names): + return value + missing: Final = next(name for name in names if not hasattr(value, name)) + raise AttributeError(f"'{type(value).__name__}' object has no attribute '{missing}'", name=missing, obj=value) + + def iter_items(value: object) -> Iterable[object]: if isinstance(value, _SupportsIter): return value @@ -329,11 +355,14 @@ def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metad Helper to get litellm metadata for spend tracking PATCH for issue where both `litellm_metadata` and `metadata` are present in the kwargs - and user_api_key values are in 'metadata'. + and user_api_key values are in 'metadata'. The proxy receive stamp rides along for the same + reason: a guardrail hook can plant `litellm_metadata` on a chat request after the proxy stamped + `metadata`, and the response timing headers anchor on that stamp. """ potential_spend_tracking_metadata_substring: Final = "user_api_key" for key, value in dict(metadata).items(): - if potential_spend_tracking_metadata_substring in key: + carry_receive_stamp = key == "litellm_received_at" and key not in litellm_metadata + if potential_spend_tracking_metadata_substring in key or carry_receive_stamp: litellm_metadata[key] = value return litellm_metadata @@ -631,11 +660,14 @@ def preserve_upstream_non_openai_attributes( for key in candidate_keys ) if all(value is None or isinstance(value, (str, int, float, bool)) for _, value in candidate_values): - for key, value in candidate_values: - setattr(model_response, key, value) + _assign_attributes(model_response, candidate_values) return - for key, value in original_chunk.model_dump(include=frozenset(candidate_keys)).items(): - setattr(model_response, key, value) + _assign_attributes(model_response, original_chunk.model_dump(include=frozenset(candidate_keys)).items()) + + +def _assign_attributes(target: object, values: Iterable[tuple[str, object]]) -> None: + for key, value in values: + setattr(target, key, value) def safe_deep_copy(data): @@ -761,6 +793,10 @@ def filter_exceptions_from_params(data: object, max_depth: int = 20) -> Any: Filtered data structure with Exception and callable objects removed, or None if the entire input was an Exception or callable """ + return _filter_exceptions_from_params(data, max_depth) + + +def _filter_exceptions_from_params(data: object, max_depth: int) -> object: if max_depth <= 0: return data @@ -782,7 +818,7 @@ def filter_exceptions_from_params(data: object, max_depth: int = 20) -> Any: if isinstance(v, Exception) or (callable(v) and not isinstance(v, type)): continue try: - filtered = filter_exceptions_from_params(v, max_depth - 1) + filtered = _filter_exceptions_from_params(v, max_depth - 1) if filtered is not None: result[k] = filtered except Exception: @@ -796,7 +832,7 @@ def filter_exceptions_from_params(data: object, max_depth: int = 20) -> Any: if isinstance(item, Exception) or (callable(item) and not isinstance(item, type)): continue try: - filtered = filter_exceptions_from_params(item, max_depth - 1) + filtered = _filter_exceptions_from_params(item, max_depth - 1) if filtered is not None: result_list.append(filtered) except Exception: @@ -892,6 +928,71 @@ class _CarriesHiddenParams(Protocol): def hidden_params(self, hidden_params: dict[str, object]) -> None: ... # mutable-ok: API requires mutation +_USAGE_BLOCK_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _usage_field(usage: object, name: str) -> object: + if isinstance(usage, Mapping): + return _USAGE_BLOCK_ADAPTER.validate_python(usage).get(name) + return getattr(usage, name, None) + + +def _cost_number(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + return float(value) + + +def _reported_cost_value(value: object) -> float | None: + """A reported cost is a number, or a breakdown whose total lives under ``total_cost``.""" + raw: Final[object] = value + if isinstance(value, Mapping): + return _cost_number(_usage_field(raw, "total_cost")) + return _cost_number(value) + + +def response_cost_in_hidden_params(response: _CarriesHiddenParams) -> float | None: + additional_headers: Final[object] = response.hidden_params.get("additional_headers") + if not isinstance(additional_headers, Mapping): + return None + return _cost_number(_USAGE_BLOCK_ADAPTER.validate_python(additional_headers).get(RESPONSE_COST_HEADER)) + + +def override_stamped_response_cost(response: _CarriesHiddenParams, cost: float) -> None: + """Replace the cost an earlier builder stamped so every reader of the hidden params agrees with the header.""" + set_response_cost_in_hidden_params(response, cost) + hidden_params: Final = response.hidden_params + hidden_params["response_cost"] = cost + + +def chunk_usage(chunk: object) -> object: + """Stream chunks reach the builders both as mappings and as ModelResponseStream objects.""" + if isinstance(chunk, Mapping): + return _USAGE_BLOCK_ADAPTER.validate_python(chunk).get("usage") + return getattr(chunk, "usage", None) + + +_MESSAGES_REPORTED_COST_PROVIDERS: Final = frozenset({"opencode_go"}) + + +def native_messages_reported_cost(usage: object, custom_llm_provider: object) -> float | None: + """Only a Messages upstream known to price its own responses is trusted, every other one keeps the price map.""" + if not isinstance(custom_llm_provider, str) or custom_llm_provider not in _MESSAGES_REPORTED_COST_PROVIDERS: + return None + return provider_reported_usage_cost(usage, custom_llm_provider) + + +def provider_reported_usage_cost(usage: object, custom_llm_provider: str | None) -> float | None: + """The request cost a provider states on its usage block, as a number or a breakdown with ``total_cost``. + + OpenCode Go reports the model's list price (peak hours and long-context tiers included) under + ``usage.opencode.list_cost``, while its ``usage.cost`` is what hit the prepaid balance, 0 inside the plan. + """ + if custom_llm_provider == "opencode_go": + return _reported_cost_value(_usage_field(_usage_field(usage, "opencode"), "list_cost")) + return _reported_cost_value(_usage_field(usage, "cost")) + + def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: float | None) -> None: """Record a provider-reported cost where the cost calculator looks before the price map.""" if cost is None: diff --git a/litellm/litellm_core_utils/cost_discount.py b/litellm/litellm_core_utils/cost_discount.py index a64ad0b20f7..bfcee7928fb 100644 --- a/litellm/litellm_core_utils/cost_discount.py +++ b/litellm/litellm_core_utils/cost_discount.py @@ -23,7 +23,7 @@ def parse_cost_discount_key(key: str) -> CostDiscountKey: def _literal_length(pattern: str) -> int: - return sum(1 for token in _GLOB_TOKEN.findall(pattern) if len(token) == 1 and token not in "*?") + return sum(1 for match in _GLOB_TOKEN.finditer(pattern) if len(match.group()) == 1 and match.group() not in "*?") def _model_names(model: str, provider_prefix: str) -> tuple[str, ...]: diff --git a/litellm/litellm_core_utils/default_encoding.py b/litellm/litellm_core_utils/default_encoding.py index c3b6a008411..970dd512a40 100644 --- a/litellm/litellm_core_utils/default_encoding.py +++ b/litellm/litellm_core_utils/default_encoding.py @@ -48,5 +48,5 @@ for attempt in range(_max_retries): # Last attempt, re-raise the exception raise # Exponential backoff with jitter to reduce collision probability - delay = _retry_delay * (2**attempt) + random.uniform(0, 0.1) + delay = _retry_delay * (2.0**attempt) + random.uniform(0, 0.1) time.sleep(delay) diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index 1dac67ecbf6..818cf1b9695 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -181,7 +181,7 @@ def _delete_nested_value_custom( def delete_nested_value( - data: dict[str, Any], + data: Mapping[str, object], path: str, depth: int = 0, max_depth: int = 20, @@ -214,7 +214,7 @@ def delete_nested_value( """ import copy - result: Final = copy.deepcopy(data) + result: Final = copy.deepcopy(data if isinstance(data, dict) else dict(data)) try: # Parse path into segments diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 1f49cc2c4bb..1c816b22ef1 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -16,7 +16,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, should_report_bug, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.secret_redaction import redact_string as redact_secret_string from litellm.types.utils import LlmProviders @@ -102,6 +102,7 @@ class ExceptionCheckers: "`inputs` tokens + `max_new_tokens` must be", "exceeds the available context size", # llama.cpp/Lemonade "exceeds the maximum number of tokens allowed", # Gemini + '"code": "context_length_exceeded"', ] for substring in known_exception_substrings: if substring in _error_str_lowercase: @@ -168,7 +169,7 @@ def get_error_message(error_obj) -> str | None: return None if hasattr(error_obj, "body"): - _error_obj_body: Final = getattr(error_obj, "body") + _error_obj_body: Final = attribute_or_none(error_obj, "body") if isinstance(_error_obj_body, dict): # OpenAI-style: {"message": "...", "type": "...", ...} if _error_obj_body.get("message"): diff --git a/litellm/litellm_core_utils/get_blog_posts.py b/litellm/litellm_core_utils/get_blog_posts.py index c8bc184f972..c3b4114d5b9 100644 --- a/litellm/litellm_core_utils/get_blog_posts.py +++ b/litellm/litellm_core_utils/get_blog_posts.py @@ -19,6 +19,7 @@ import httpx from pydantic import ConfigDict from litellm import verbose_logger +from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.types.llms.base import LiteLLMBaseModel BLOG_POSTS_TTL_SECONDS: Final[int] = 3600 # 1 hour @@ -129,7 +130,7 @@ class GetBlogPosts: Uses the in-process cache if within BLOG_POSTS_TTL_SECONDS. Fetches from ``url`` otherwise, falling back to local backup on failure. """ - if os.getenv("LITELLM_LOCAL_BLOG_POSTS", "").lower() == "true": + if is_offline_mode() or os.getenv("LITELLM_LOCAL_BLOG_POSTS", "").lower() == "true": return cls.load_local_blog_posts() now: Final = time.time() diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 85557090ae9..ca8e2af18e6 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -76,6 +76,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "itpm", "otpm", "use_xai_oauth", + "chatgpt_auth_file", "github_copilot_auth_type", "github_copilot_host", "github_copilot_user_session", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index a67c1bb7df8..19c1d5c414c 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -158,6 +158,14 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None +def resolve_model_and_provider_for_metadata(model: str, custom_llm_provider: str | None = None) -> tuple[str, str]: + declared: Final = declared_authenticating_provider(model, custom_llm_provider) + if declared is not None: + return model.removeprefix(f"{declared}/"), declared + resolved_model, resolved_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + return resolved_model, resolved_provider + + def inferred_provider(model: str | None) -> str | None: if not model: return None @@ -842,7 +850,9 @@ def _get_openai_compatible_provider_info( api_base, dynamic_api_key, custom_llm_provider, - ) = litellm.ChatGPTConfig().get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider) + ) = litellm.ChatGPTConfig().get_openai_compatible_provider_info( + model, api_base, api_key, custom_llm_provider, litellm_params + ) elif custom_llm_provider == "novita": api_base = api_base or get_secret("NOVITA_API_BASE") or "https://api.novita.ai/v3/openai" dynamic_api_key = api_key or get_secret_str("NOVITA_API_KEY") @@ -902,6 +912,10 @@ def _get_openai_compatible_provider_info( api_base, dynamic_api_key, ) = litellm.MorphChatConfig().get_openai_compatible_provider_info(api_base, api_key) + elif custom_llm_provider == "opencode_go": + from litellm.llms.opencode.common_utils import go_api_base, go_api_key + + return model, custom_llm_provider, go_api_key(api_key), go_api_base(api_base, model) elif custom_llm_provider == "lambda_ai": ( api_base, @@ -972,3 +986,9 @@ def _get_openai_compatible_provider_info( if dynamic_api_key is None and api_key is not None: dynamic_api_key = api_key return model, custom_llm_provider, dynamic_api_key, api_base + + +def get_provider_affinity_header(api_base: str | None) -> str | None: + from litellm.llms.opencode.transformation import OpenCodeProviderAdapter + + return OpenCodeProviderAdapter().get_provider_affinity_header(api_base) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 7b3835dc1ac..7fc6eab6c60 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -39,9 +39,11 @@ from litellm.constants import ( from litellm.litellm_core_utils.fallback_generalizations import ( set_fallback_generalizations, ) +from litellm.litellm_core_utils.offline_mode import is_offline_mode FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) +_DECODED_JSON: Final = TypeAdapter(object) _BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) @@ -313,7 +315,7 @@ def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) - def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _FetchAttemptOutcome: try: - parsed: Final[object] = json.loads(content) + parsed: Final = _DECODED_JSON.validate_python(json.loads(content)) except ValueError as e: return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}") if not isinstance(parsed, dict): @@ -382,7 +384,8 @@ def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemp ) if response.is_error: return ModelCostMapReloadUnavailable(reason=f"HTTP {response.status_code} from {url}") - return _classify_fetch_content(response.content, url, response.headers.get("etag")) + etag_values: Final = response.headers.get_list("etag") + return _classify_fetch_content(response.content, url, ", ".join(etag_values) if etag_values else None) def _next_retry_wait( @@ -456,7 +459,7 @@ async def refetch_model_cost_map( on failure it returns ``ModelCostMapReloadUnavailable`` so callers keep the map they already have. """ - if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": + if is_offline_mode() or os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": _cost_map_source_info.loaded_at = datetime.now(timezone.utc) _cost_map_source_info.source = "local" _cost_map_source_info.url = None @@ -740,7 +743,7 @@ def get_model_cost_map( _cost_map_source_info.loaded_at = datetime.now(timezone.utc) # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. - if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): + if is_offline_mode() or os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True diff --git a/litellm/litellm_core_utils/get_provider_specific_headers.py b/litellm/litellm_core_utils/get_provider_specific_headers.py index bc325c3c702..62dfcb073e5 100644 --- a/litellm/litellm_core_utils/get_provider_specific_headers.py +++ b/litellm/litellm_core_utils/get_provider_specific_headers.py @@ -1,10 +1,23 @@ from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter + from litellm.types.utils import ProviderSpecificHeader +_PROVIDER_HEADERS: Final[TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None]] = TypeAdapter( + ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None +) + class ProviderSpecificHeaderUtils: + @staticmethod + def forwarded_authorization(provider_specific_header: object, custom_llm_provider: str | None) -> str | None: + headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( + _PROVIDER_HEADERS.validate_python(provider_specific_header), custom_llm_provider + ) + return next((value for key, value in headers.items() if key.lower() == "authorization"), None) + @staticmethod def get_provider_specific_headers( provider_specific_header: ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None, diff --git a/litellm/litellm_core_utils/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index f8d44d44cdc..a74337d2549 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -1,4 +1,5 @@ from collections.abc import Iterator, Mapping, MutableMapping +from types import MappingProxyType from typing import Final, Protocol, cast from pydantic import BaseModel @@ -52,7 +53,7 @@ class HiddenParamsModelView(MutableMapping[str, object]): def stored_attribute(obj: object, name: str) -> object: if not isinstance(obj, BaseModel): return getattr(obj, name, None) - values: Final[Mapping[str, object]] = vars(obj) + values: Final[Mapping[str, object]] = MappingProxyType[str, object](vars(obj)) if name in values: return values[name] extras: Final[Mapping[str, object]] = obj.__pydantic_extra__ or {} diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index 5a533054f0b..c3769686e45 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -28,7 +28,7 @@ def inherit_message_logging_privacy(disabled: bool) -> Generator[None]: def iter_client_callback_metadata_dicts( - kwargs: dict[str, Any], + kwargs: Mapping[str, object], ) -> Iterator[tuple[str, dict[str, Any]]]: litellm_params: Final = kwargs.get("litellm_params") if isinstance(litellm_params, dict): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 473666e5ba1..9e0461e951d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -12,7 +12,7 @@ import subprocess import sys import time import traceback -from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime as dt_object from functools import lru_cache @@ -77,6 +77,8 @@ from litellm.litellm_core_utils.classifier_logging import ( from litellm.litellm_core_utils.core_helpers import ( get_provider_response_headers_from_hidden_params, is_expected_client_error, + native_messages_reported_cost, + provider_reported_usage_cost, proxy_stamped_kubernetes_pod_routing, proxy_stamped_used_client_oauth_token, reconstruct_model_name, @@ -138,6 +140,7 @@ from litellm.types.interactions import ( InteractionsAPIResponse, InteractionsAPIStreamingResponse, ) +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.openai import ( AllMessageValues, Batch, @@ -160,6 +163,7 @@ from litellm.types.utils import ( CachingDetails, CallTypes, CostBreakdown, + CostPerToken, CustomPricingLiteLLMParams, DynamicPromptManagementParamLiteral, EmbeddingResponse, @@ -554,6 +558,7 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) +_REDACTED_RESULT: Final = TypeAdapter(object) def _provider_response_id(source: object) -> str | None: @@ -574,6 +579,27 @@ class UncachedPrice: cost_breakdown: CostBreakdown | None +@dataclass(frozen=True, slots=True) +class BilledRoundsSpend: + usage: Usage + cost: float + cost_breakdown: CostBreakdown | None + + +@dataclass(frozen=True, slots=True) +class BilledRoundsDispatch: + """A parked success log for a stream bridged through chat completions that billed several rounds, and the spend + its failure log bills when a post-call hook rejects the response, as it does for a rejected one-round bridged + stream.""" + + dispatch: Coroutine[object, object, None] + spend: BilledRoundsSpend + + +def _fixed_margin(cost_breakdown: CostBreakdown | None) -> float: + return 0.0 if cost_breakdown is None else cost_breakdown.get("margin_fixed_amount", 0.0) + + def _timestamp_seconds(moment: object) -> float | None: if isinstance(moment, datetime.datetime): return moment.timestamp() @@ -582,6 +608,10 @@ def _timestamp_seconds(moment: object) -> float | None: return None +def same_attempt_marker(first: object, second: object) -> bool: + return first is second or first == second + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -595,6 +625,9 @@ class Logging(LiteLLMLoggingBaseClass): prometheusLogger, \ slack_app custom_pricing: bool = False + custom_cost_per_token: CostPerToken | None = None + custom_cost_per_second: float | None = None + replaced_global_callback_keys: frozenset[str] = frozenset() stream_options = None litellm_request_debug: bool = False streamed_anthropic_message_id: str | None = None @@ -602,6 +635,7 @@ class Logging(LiteLLMLoggingBaseClass): classifier_input: Mapping[str, JsonValue] | None = None baseline_cache_context: "BaselineCacheContext | None" = None baseline_observation: "CapturedBaselineObservation | None" = None + _on_failure_payload_built: "Callable[[], None] | None" = None @property def _defer_async_logging(self) -> bool: @@ -649,11 +683,11 @@ class Logging(LiteLLMLoggingBaseClass): litellm_call_id: str, function_id: str, litellm_trace_id: str | None = None, - dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, - dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None, - dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, + dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, applied_guardrails: list[str] | None = None, kwargs: dict | None = None, log_raw_request_response: bool = False, @@ -698,6 +732,7 @@ class Logging(LiteLLMLoggingBaseClass): self._pre_call_session_id: str = session_id_var.get() _sid: Final = kwargs.get("litellm_session_id") if kwargs else None self.litellm_session_id: str = str(_sid) if _sid else "" + self.stream_usage_injected: Final = kwargs is not None and kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True # supports_correlation_logging is False for calls originating from the # sync client entry point (wrapper() in utils.py): a plain OS thread # has no per-call context isolation the way an asyncio Task does, and @@ -726,15 +761,17 @@ class Logging(LiteLLMLoggingBaseClass): self.raw_request_only = raw_request_only # Initialize dynamic callbacks - self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks + self.dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = dynamic_input_callbacks self.dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_success_callbacks ) self.dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_async_success_callbacks ) - self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks - self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = ( + self.dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( + dynamic_failure_callbacks + ) + self.dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_async_failure_callbacks ) @@ -826,6 +863,7 @@ class Logging(LiteLLMLoggingBaseClass): "litellm_call_id": litellm_call_id, "start_time": start_time, "first_api_call_start_time": None, + "current_attempt_api_call_start_time": None, "litellm_params": litellm_params, } new_call.streaming_chunks = [] @@ -851,10 +889,32 @@ class Logging(LiteLLMLoggingBaseClass): self.dynamic_async_failure_callbacks, callback ) + def replace_global_callback(self, callback: CustomLogger) -> None: + self.add_dynamic_callback(callback) + self.replaced_global_callback_keys = self.replaced_global_callback_keys | { + litellm.logging_callback_manager.get_custom_logger_key(callback) + } + + def add_dynamic_function_callbacks( + self, + *, + async_success: Callable[..., object], + failure: Callable[..., object], + async_failure: Callable[..., object], + ) -> None: + self.dynamic_async_success_callbacks = self._with_dynamic_callback( + self.dynamic_async_success_callbacks, async_success + ) + self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, failure) + self.dynamic_async_failure_callbacks = self._with_dynamic_callback( + self.dynamic_async_failure_callbacks, async_failure + ) + @staticmethod def _with_dynamic_callback( - callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger - ) -> list[str | Callable | CustomLogger]: + callbacks: Sequence[str | Callable[..., object] | CustomLogger] | None, + callback: Callable[..., object] | CustomLogger, + ) -> list[str | Callable[..., object] | CustomLogger]: existing: Final = tuple(callbacks or ()) return [*existing, *(() if callback in existing else (callback,))] @@ -891,9 +951,9 @@ class Logging(LiteLLMLoggingBaseClass): def _process_dynamic_callback_list( self, - callback_list: list[str | Callable | CustomLogger] | None, + callback_list: list[str | Callable[..., object] | CustomLogger] | None, dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"], - ) -> list[str | Callable | CustomLogger] | None: + ) -> list[str | Callable[..., object] | CustomLogger] | None: """ Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks @@ -905,7 +965,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback_list is None: return None - processed_list: Final[list[str | Callable | CustomLogger]] = [] + processed_list: Final[list[str | Callable[..., object] | CustomLogger]] = [] for callback in callback_list: if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks: for callback_instance in self._resolve_dynamic_callback_string(callback): @@ -1114,7 +1174,7 @@ class Logging(LiteLLMLoggingBaseClass): automatically extracts metadata/litellm_metadata from kwargs, so callers don't need to manually plumb them into litellm_params. """ - base_litellm_params: Final[dict[str, Any]] = {} + base_litellm_params: Final[dict[str, object]] = {} if isinstance(kwargs.get("metadata"), dict): base_litellm_params["metadata"] = kwargs["metadata"].copy() @@ -1131,10 +1191,11 @@ class Logging(LiteLLMLoggingBaseClass): lp_metadata: Final = litellm_params.pop("metadata", None) base_litellm_params.update(litellm_params) if lp_metadata and isinstance(lp_metadata, dict): - base_litellm_params.setdefault("metadata", {}) - for k, v in lp_metadata.items(): - if k not in base_litellm_params["metadata"]: - base_litellm_params["metadata"][k] = v + merged_metadata: Final = base_litellm_params.setdefault("metadata", {}) + if isinstance(merged_metadata, dict): + for k, v in lp_metadata.items(): + if k not in merged_metadata: + merged_metadata[k] = v self.update_environment_variables( litellm_params=base_litellm_params, @@ -1554,7 +1615,9 @@ class Logging(LiteLLMLoggingBaseClass): # A proxy-gate failure (auth, rate limit, client disconnect) fires pre_call for its callbacks # without any provider handoff, so it must not stamp the handoff instants. if LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL not in self.model_call_details: - self.model_call_details["api_call_start_time"] = datetime.datetime.now() + api_call_start_time: Final = datetime.datetime.now() + self.model_call_details["api_call_start_time"] = api_call_start_time + self.model_call_details["current_attempt_api_call_start_time"] = api_call_start_time # Set-once first provider-handoff instant. api_call_start_time # is overwritten on every retry, so it can't measure one-time # preprocessing; pinning the first attempt excludes retry loops @@ -1562,11 +1625,12 @@ class Logging(LiteLLMLoggingBaseClass): # litellm_params["metadata"] (caller request metadata, typed # Dict[str, str], echoed downstream; a datetime breaks it). if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = self.model_call_details[ - "api_call_start_time" - ] + self.model_call_details["first_api_call_start_time"] = api_call_start_time # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) + callbacks: Final = ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) for callback in callbacks: try: if callback == "supabase" and supabaseClient is not None: @@ -1733,13 +1797,16 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.exception( "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - original_response = redact_message_input_output_from_logging( + redact_message_input_output_from_logging( model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=original_response, ) # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) + callbacks: Final = ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) for callback in callbacks: try: if callback == "sentry" and add_breadcrumb: @@ -1960,10 +2027,9 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details.pop("response_cost_failure_debug_information", None) - if cache_hit is None: - cache_hit = self.model_call_details.get("cache_hit", False) + resolved_cache_hit: Final = self.model_call_details.get("cache_hit", False) if cache_hit is None else cache_hit - if cache_hit is True: + if resolved_cache_hit is True: return 0.0 if is_unbilled_non_inference_call( @@ -2015,7 +2081,7 @@ class Logging(LiteLLMLoggingBaseClass): model_value: Final = litellm_model_name or self.model cost_cache_hit_value: Final = ( - self.model_call_details.get("cache_hit", False) if cache_hit is None else cache_hit + self.model_call_details.get("cache_hit", False) if resolved_cache_hit is None else resolved_cache_hit ) cost_cache_hit: Final[bool | None] = cast( # cast-ok: callback metadata is caller-provided bool | None, cost_cache_hit_value @@ -2036,6 +2102,8 @@ class Logging(LiteLLMLoggingBaseClass): "call_type": self.call_type, "optional_params": self.optional_params, "custom_pricing": custom_pricing, + "custom_cost_per_token": self.custom_cost_per_token if custom_pricing else None, + "custom_cost_per_second": self.custom_cost_per_second if custom_pricing else None, "prompt": prompt, "standard_built_in_tools_params": self.standard_built_in_tools_params, "router_model_id": router_model_id, @@ -2388,6 +2456,31 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def record_failure_after_earlier_rounds( + self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float + ) -> None: + """Bill a failed round on the failure log after the rounds of its request before it were billed on the success + log. + + The fixed margin is a per-request fee, so it is charged once: the failed round's margin is dropped when the + earlier rounds already paid one.""" + from litellm.responses.mcp.round_billing import breakdown_less_repeated_fixed_margin + + failed_breakdown: Final = None if self.cost_breakdown is earlier.cost_breakdown else self.cost_breakdown + repeated_fixed_margin: Final = ( + 0.0 + if failed_breakdown is None + else min(_fixed_margin(earlier.cost_breakdown), _fixed_margin(failed_breakdown), earlier.cost, failed_cost) + ) + self.model_call_details["combined_usage_object"] = failed_usage + self.model_call_details["response_cost"] = failed_cost - repeated_fixed_margin + self.cost_breakdown = breakdown_less_repeated_fixed_margin(failed_breakdown, repeated_fixed_margin) + + def record_billed_rounds_for_failure(self, spend: BilledRoundsSpend) -> None: + """Bill every round of a fully streamed request on the failure log when a post-call hook rejects it.""" + self.record_partial_usage_for_failure(spend.usage, spend.cost) + self.cost_breakdown = spend.cost_breakdown + def attempt_scoped_copy(self) -> "Logging": """A copy with its own model_call_details, so logging one attempt of a request through it leaves the state the request's other attempts share untouched.""" @@ -2395,12 +2488,39 @@ class Logging(LiteLLMLoggingBaseClass): attempt.model_call_details = dict(self.model_call_details) return attempt + def dispatch_attempt_failure_handlers( + self, exception: Exception, traceback_exception: str + ) -> Coroutine[None, None, None]: + shared_payload: Final = self.model_call_details.get("standard_logging_object") + attempt_marker: Final = self.model_call_details.get("current_attempt_api_call_start_time") + attempt: Final = self.attempt_scoped_copy() + + def publish() -> None: + attempt_payload: Final = attempt.model_call_details.get("standard_logging_object") + current_marker: Final = self.model_call_details.get("current_attempt_api_call_start_time") + if ( + attempt_payload is not None + and self.model_call_details.get("standard_logging_object") is shared_payload + and same_attempt_marker(current_marker, attempt_marker) + ): + self.model_call_details["standard_logging_object"] = attempt_payload + + attempt._on_failure_payload_built = publish + + async def dispatch() -> None: + await attempt.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) + + return dispatch() + def carry_failed_attempt_usage(self) -> None: """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its own usage while the request's failure log still bills every attempt.""" - usage: Final = self.model_call_details.pop("combined_usage_object", None) - response_cost: Final = self.model_call_details.pop("response_cost", None) + usage: Final = self._model_call_details_view.get("combined_usage_object") + response_cost: Final = self._model_call_details_view.get("response_cost") + self.model_call_details.pop("combined_usage_object", None) + self.model_call_details.pop("response_cost", None) if isinstance(usage, Usage): + self.model_call_details["current_attempt_api_call_start_time"] = None self.model_call_details["failed_attempts_usage"] = ( *self.model_call_details.get("failed_attempts_usage", ()), (usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0), @@ -2514,6 +2634,10 @@ class Logging(LiteLLMLoggingBaseClass): """ return self.model_call_details[f"has_logged_{event_type}"] = True + if event_type == "async_failure": + self.model_call_details["async_failure_logged_attempt"] = self.model_call_details.get( + "current_attempt_api_call_start_time" + ) return def has_run_logging( @@ -3183,7 +3307,7 @@ class Logging(LiteLLMLoggingBaseClass): 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: dict[str, object] # 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 @@ -3204,9 +3328,11 @@ class Logging(LiteLLMLoggingBaseClass): ) ## REDACT MESSAGES ## - result = redact_message_input_output_from_logging( - model_call_details=call_details, - result=result, + result = _REDACTED_RESULT.validate_python( + redact_message_input_output_from_logging( + model_call_details=call_details, + result=result, + ) ) ## LOGGING HOOK ## hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks))) @@ -3472,7 +3598,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose("S3Logger Logger: Got Stream Event - Completed Stream Response") s3Logger.log_event( kwargs=call_details, - response_obj=call_details["complete_streaming_response"], + response_obj=call_details_view["complete_streaming_response"], start_time=start_time, end_time=end_time, print_verbose=print_verbose, @@ -3795,9 +3921,11 @@ class Logging(LiteLLMLoggingBaseClass): self._get_completed_success_callbacks(asynchronous=True) if not self.stream else callbacks ) - result = redact_message_input_output_from_logging( - model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), - result=result, + result = _REDACTED_RESULT.validate_python( + redact_message_input_output_from_logging( + model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), + result=result, + ) ) ## LOGGING HOOK ## @@ -4142,9 +4270,11 @@ class Logging(LiteLLMLoggingBaseClass): result: object = None # result sent to all loggers, init this to None incase it's not created - result = redact_message_input_output_from_logging( - model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), - result=result, + result = _REDACTED_RESULT.validate_python( + redact_message_input_output_from_logging( + model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), + result=result, + ) ) self.mark_logging_complete(event_type="sync_failure") for callback in callbacks: @@ -4331,7 +4461,8 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, ) - + if self._on_failure_payload_built is not None: + self._on_failure_payload_built() callbacks: Final = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_async_failure_callbacks, global_callbacks=litellm._async_failure_callback, @@ -4516,14 +4647,26 @@ class Logging(LiteLLMLoggingBaseClass): _filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks) return len(_filtered_failure_callbacks) > 0 + def _without_replaced_global_callbacks( + self, global_callbacks: Iterable[str | Callable[..., object] | CustomLogger] + ) -> tuple[str | Callable[..., object] | CustomLogger, ...]: + if not self.replaced_global_callback_keys: + return tuple(global_callbacks) + return tuple( + litellm.logging_callback_manager.without_custom_logger_keys( + global_callbacks, self.replaced_global_callback_keys + ) + ) + def get_combined_callback_list( self, dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None, global_callbacks: list[str | Callable[..., object] | CustomLogger], ) -> list[str | Callable[..., object] | CustomLogger]: + kept_global_callbacks: Final = self._without_replaced_global_callbacks(global_callbacks) if dynamic_success_callbacks is None: - return list(global_callbacks) - return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks)) + return list(kept_global_callbacks) + return list(dict.fromkeys((*dynamic_success_callbacks, *kept_global_callbacks))) def _remove_internal_litellm_callbacks(self, callbacks: list) -> list: """ @@ -4615,7 +4758,12 @@ class Logging(LiteLLMLoggingBaseClass): ): ## return unified Usage object if isinstance(result.response.usage, ResponseAPIUsage): - set_response_cost_in_hidden_params(result.response, result.response.usage.cost) + set_response_cost_in_hidden_params( + result.response, + provider_reported_usage_cost( + result.response.usage, self.model_call_details.get("custom_llm_provider") + ), + ) transformed_usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage( result.response.usage ) @@ -4681,8 +4829,11 @@ class Logging(LiteLLMLoggingBaseClass): bridge is replaced by the ModelResponse its spend row is built from, anything else is priced as is.""" served: Final[object] = result - if self.bridged_model_response is None or not isinstance(result, dict): - return result + call_type: Final[object] = getattr(self, "call_type", None) + if not isinstance(result, dict): + return served + if self.bridged_model_response is None and call_type != CallTypes.anthropic_messages.value: + return served return self._anthropic_messages_logged_response(result=served) def _anthropic_messages_logged_response(self, result: object) -> ModelResponse: @@ -4725,7 +4876,13 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(result, ResponsesAPIResponse): return self._translate_responses_api_response_to_model_response(result) - provider_response_id: Final = _provider_response_id(result) + raw_body: Final[object] = result + served_body: Final = _LITELLM_PARAMS.validate_python(raw_body) if isinstance(result, dict) else None + reported_cost: Final = native_messages_reported_cost( + served_body.get("usage") if served_body is not None else None, + self.model_call_details.get("custom_llm_provider"), + ) + provider_response_id: Final = _provider_response_id(raw_body) bridged: Final = self.bridged_model_response if bridged is not None and isinstance(result, dict): return bridged if provider_response_id is None else bridged.model_copy(update={"id": provider_response_id}) @@ -4750,7 +4907,7 @@ class Logging(LiteLLMLoggingBaseClass): else: from litellm.types.llms.anthropic import AnthropicResponse - pydantic_result: Final = AnthropicResponse.model_validate(result) + pydantic_result: Final = AnthropicResponse.model_validate({"stop_sequence": None, **(served_body or {})}) import httpx result = litellm.AnthropicConfig().transform_parsed_response( @@ -4763,6 +4920,7 @@ class Logging(LiteLLMLoggingBaseClass): json_mode=None, speed=self.optional_params.get("speed") if self.optional_params else None, ) + set_response_cost_in_hidden_params(result, reported_cost) return result def _translate_responses_api_response_to_model_response(self, result: ResponsesAPIResponse) -> ModelResponse: @@ -4820,7 +4978,7 @@ class Logging(LiteLLMLoggingBaseClass): """ import httpx - httpx_response: Final = self.model_call_details.get("httpx_response", None) + httpx_response: Final[Response | None] = self.model_call_details.get("httpx_response", None) if httpx_response is None: raise ValueError("Google GenAI Generate Content: httpx_response is None") dict_result: Final = httpx_response.json() @@ -4850,7 +5008,7 @@ class Logging(LiteLLMLoggingBaseClass): The response object with usage added if available """ # Get usage from model_call_details (set by asend_message) - usage: Final = self.model_call_details.get("usage") + usage: Final[Usage | None] = self.model_call_details.get("usage") if usage is None: return result @@ -5683,7 +5841,9 @@ def build_custom_logger_compatible_class( return callback # Get global BitBucket config - bitbucket_config: Final = getattr(litellm, "global_bitbucket_config", None) + bitbucket_config: Final = ( + litellm.global_bitbucket_config if hasattr(litellm, "global_bitbucket_config") else None + ) if bitbucket_config is None: raise ValueError("BitBucket configuration not found. Please set litellm.global_bitbucket_config first.") diff --git a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py index 9bcc2b1743c..3040e36ec5f 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py @@ -5,30 +5,42 @@ Shared by provider cost calculators (e.g. Dashscope) and the proxy budget reservation logic so neither has to depend on the other. """ -from typing import Final +from collections.abc import Mapping, Sequence +from typing import Final, SupportsFloat + +from pydantic import TypeAdapter, ValidationError + +_TIER_RANGE: Final = TypeAdapter(tuple[float, float]) -def _coerce_cost_per_token(value: float | str | None) -> float: +def _coerce_cost_per_token(value: object) -> float: """ Coerce a per-token cost into a float. Model cost values loaded from YAML config may arrive as strings (e.g. scientific notation like "4e-07"), which would break arithmetic. """ - if value is None: - return 0.0 if isinstance(value, str): try: return float(value) except ValueError: return 0.0 - return float(value) + if isinstance(value, SupportsFloat): + return float(value) + return 0.0 + + +def _tier_range(tier: Mapping[str, object]) -> tuple[float, float] | None: + try: + return _TIER_RANGE.validate_python(tier.get("range")) + except ValidationError: + return None def select_tier_for_input( - tiered_pricing: list[dict], + tiered_pricing: Sequence[Mapping[str, object]], input_tokens: int, -) -> dict | None: +) -> Mapping[str, object] | None: """ Select the pricing tier for a request based on its total input token count. @@ -44,19 +56,21 @@ def select_tier_for_input( if not tiered_pricing or input_tokens <= 0: return None - sorted_tiers: Final = sorted(tiered_pricing, key=lambda t: t.get("range", [0, 0])[0]) - valid_tiers: Final = [tier for tier in sorted_tiers if len(tier.get("range", [])) == 2] - if not valid_tiers: + ranged_tiers: Final = sorted( + ((bounds, tier) for tier in tiered_pricing if (bounds := _tier_range(tier)) is not None), + key=lambda ranged: ranged[0][0], + ) + if not ranged_tiers: return None - matching: Final = [tier for tier in valid_tiers if tier["range"][0] < input_tokens <= tier["range"][1]] + matching: Final = tuple(tier for (start, end), tier in ranged_tiers if start < input_tokens <= end) if matching: return matching[0] - return valid_tiers[-1] + return ranged_tiers[-1][1] def tier_rate( - tier: dict, + tier: Mapping[str, object], cost_key: str, fallback_cost_key: str | None = None, ) -> float: @@ -68,4 +82,6 @@ def tier_rate( primary: Final = tier.get(cost_key) if primary is not None: return _coerce_cost_per_token(primary) + if fallback_cost_key is None: + return 0.0 return _coerce_cost_per_token(tier.get(fallback_cost_key, 0)) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 04510c20db0..045869e4547 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -450,7 +450,7 @@ class StandardBuiltInToolCostTracking: def _get_code_interpreter_cost( model_info: ModelInfo | None, custom_llm_provider: str | None, - standard_built_in_tools_params: StandardBuiltInToolsParams, + standard_built_in_tools_params: Mapping[str, object], ) -> float: """Calculate code interpreter cost.""" code_interpreter_sessions: Final = standard_built_in_tools_params.get("code_interpreter_sessions", None) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index a874459798d..3b49f239afe 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -3,7 +3,7 @@ import re from collections.abc import Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field, replace from datetime import datetime, timezone, tzinfo from types import MappingProxyType from typing import Final, Literal, TypedDict, cast @@ -52,6 +52,22 @@ _VALID_DATA_RESIDENCIES: Final = frozenset(r.value for r in DataResidency) _DEPLOYMENT_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields) +_BASE_CACHE_WRITE_RATE_KEY: Final = "cache_creation_input_token_cost" + +_CACHE_WRITE_RATE_KEYS: Final[Mapping[str, str]] = MappingProxyType( + { + "5m": _BASE_CACHE_WRITE_RATE_KEY, + "30m": _BASE_CACHE_WRITE_RATE_KEY, + "1hr": "cache_creation_input_token_cost_above_1hr", + "12hr": "cache_creation_input_token_cost_above_12hr", + } +) + +_TTL_PRICED_CACHE_WRITE_RATE_KEYS: Final[Mapping[str, str]] = MappingProxyType( + {ttl: key for ttl, key in _CACHE_WRITE_RATE_KEYS.items() if key != _BASE_CACHE_WRITE_RATE_KEY} +) + + _IMAGE_TOKEN_RATE_KEYS: Final[tuple[str, ...]] = ( "input_cost_per_token", "output_cost_per_token", @@ -88,7 +104,8 @@ _BATCH_RATE_PREFIXES: Final = ( "cache_read_input_token_cost", "cache_creation_input_token_cost", ) -_BATCH_TIER_KEY: Final = re.compile(rf"^({'|'.join(_BATCH_RATE_PREFIXES)})_above_(\d+k?)_tokens{_BATCH_KEY_SUFFIX}$") +_BATCH_TIER_PREFIXES: Final = (*_BATCH_RATE_PREFIXES, *_TTL_PRICED_CACHE_WRITE_RATE_KEYS.values()) +_BATCH_TIER_KEY: Final = re.compile(rf"^({'|'.join(_BATCH_TIER_PREFIXES)})_above_(\d+k?)_tokens{_BATCH_KEY_SUFFIX}$") _NON_STANDARD_THRESHOLD_SUFFIXES: Final = (*_SERVICE_TIER_SUFFIXES, _BATCH_KEY_SUFFIX) @@ -298,6 +315,7 @@ class BatchCostRates: output: float | None cache_read: float | None cache_creation: float | None + cache_creation_rates_by_ttl: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({})) def _batch_rate(model_info: ModelInfo, key: str) -> float | None: @@ -353,10 +371,31 @@ def get_batch_cost_rates(model_info: ModelInfo, usage: Usage, custom_llm_provide output=_batch_rate_for_prefix(model_info, "output_cost_per_token", usage, inclusive), cache_read=_batch_rate_for_prefix(model_info, "cache_read_input_token_cost", usage, inclusive), cache_creation=_batch_rate_for_prefix(model_info, "cache_creation_input_token_cost", usage, inclusive), + cache_creation_rates_by_ttl=MappingProxyType( + { + ttl: rate + for ttl, key in _TTL_PRICED_CACHE_WRITE_RATE_KEYS.items() + if (rate := _batch_rate_for_prefix(model_info, key, usage, inclusive)) is not None + } + ), ) -def _select_priced_tier(model_info: ModelInfo, usage: Usage) -> dict[str, object] | None: +@dataclass(frozen=True, slots=True) +class TokenRates: + input_rate: float + output_rate: float + cache_read_rate: float + cache_creation_rate: float + reasoning_rate: float | None + cache_creation_rates_by_ttl: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({})) + + @property + def billed_reasoning_rate(self) -> float: + return self.output_rate if self.reasoning_rate is None else self.reasoning_rate + + +def _select_priced_tier(model_info: ModelInfo, usage: Usage) -> Mapping[str, object] | None: tiered_pricing: Final = model_info.get("tiered_pricing") if not isinstance(tiered_pricing, list) or not tiered_pricing: return None @@ -376,7 +415,7 @@ def _get_tiered_reasoning_rate(model_info: ModelInfo, usage: Usage) -> float | N return tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token") -def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, float, float, float, float] | None: +def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[TokenRates, Mapping[str, object]] | None: """ Resolve the base rates from a model's ``tiered_pricing`` table, if it has one. @@ -391,22 +430,20 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, if tier is None: return None - cache_creation_cost: Final = tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") completion_cost: Final = ( tier_rate(tier, "output_cost_per_token") if "output_cost_per_token" in tier else get_cost_per_unit(model_info, "output_cost_per_token") or 0.0 ) - return ( - tier_rate(tier, "input_cost_per_token"), - completion_cost, - cache_creation_cost, - tier_rate(tier, "cache_creation_input_token_cost_above_1hr") - if "cache_creation_input_token_cost_above_1hr" in tier - and tier["cache_creation_input_token_cost_above_1hr"] is not None - else cache_creation_cost, - tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate: Final = tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") + rates: Final = TokenRates( + input_rate=tier_rate(tier, "input_cost_per_token"), + output_rate=completion_cost, + cache_read_rate=tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate=cache_creation_rate, + reasoning_rate=None, ) + return rates, tier def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_time: datetime | None = None) -> bool: @@ -532,19 +569,6 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = return False -@dataclass(frozen=True, slots=True) -class TokenRates: - input_rate: float - output_rate: float - cache_read_rate: float - cache_creation_rate: float - reasoning_rate: float | None - - @property - def billed_reasoning_rate(self) -> float: - return self.output_rate if self.reasoning_rate is None else self.reasoning_rate - - def resolve_explicit_cache_read_cost_rate( model_info: ModelInfo, usage: Usage, @@ -649,7 +673,8 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None, if off_peak is None: return rates off_peak_reasoning_rate: Final = _parse_off_peak_rate(off_peak.get("output_cost_per_reasoning_token")) - return TokenRates( + return replace( + rates, input_rate=_off_peak_rate(off_peak, "input_cost_per_token", rates.input_rate), output_rate=_off_peak_rate(off_peak, "output_cost_per_token", rates.output_rate), cache_read_rate=_off_peak_rate(off_peak, "cache_read_input_token_cost", rates.cache_read_rate), @@ -661,32 +686,28 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None, def _apply_off_peak_to_base_costs( model_info: ModelInfo, current_time: datetime | None, - base_costs: tuple[float, float, float, float | None, float], -) -> tuple[float, float, float, float, float]: - """Apply off-peak rates to an already-resolved set of base costs, whichever pricing path - produced them. off_peak_pricing has no field for the one-hour cache-creation rate, so a - present one passes through untouched and an absent one resolves to the applied - cache-creation rate. Reasoning is left to _resolve_billed_reasoning_rate. + base_rates: TokenRates, + tier: Mapping[str, object], +) -> TokenRates: + """Apply off-peak rates to the base costs resolved from a pricing tier. off_peak_pricing has no + field for the long-TTL cache-creation rates, so a tier that declares one keeps it, and a tier + that does not bills that TTL at the off-peak-adjusted write rate. Reasoning is left to + _resolve_billed_reasoning_rate. """ - prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs - rates: Final = apply_off_peak_pricing( - model_info, - current_time, - TokenRates( - input_rate=prompt, - output_rate=completion, - cache_read_rate=cache_read, - cache_creation_rate=cache_creation, - reasoning_rate=None, + rates: Final = apply_off_peak_pricing(model_info, current_time, base_rates) + return replace( + rates, + cache_creation_rates_by_ttl=MappingProxyType( + { + ttl: rates.cache_creation_rate if tier.get(key) is None else tier_rate(tier, key) + for ttl, key in _TTL_PRICED_CACHE_WRITE_RATE_KEYS.items() + } ), ) - return ( - rates.input_rate, - rates.output_rate, - rates.cache_creation_rate, - rates.cache_creation_rate if cache_creation_above_1hr is None else cache_creation_above_1hr, - rates.cache_read_rate, - ) + + +def _first_present_key(model_info: ModelInfo, preferred_key: str, fallback_key: str) -> str: + return preferred_key if model_info.get(preferred_key) is not None else fallback_key def _get_token_base_cost( @@ -696,7 +717,7 @@ def _get_token_base_cost( current_time: datetime | None = None, *, threshold_is_inclusive: bool = False, -) -> tuple[float, float, float, float, float]: +) -> TokenRates: """ Return prompt cost, completion cost, and cache costs for a given model and usage. @@ -709,15 +730,12 @@ def _get_token_base_cost( An absent cache-creation or cache-read rate always resolves to the resolved input rate, the way the tiered table and custom deployment pricing already do, since a provider that publishes no cache price bills cached tokens as ordinary input. An - absent 1h write rate resolves to the cache-creation rate, off-peak included. An + absent long-TTL write rate resolves to the cache-creation rate, off-peak included. An explicit 0.0 stays a real price for all of them. - - Returns: - Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost) """ tiered_base_costs: Final = _get_tiered_base_costs(model_info=model_info, usage=usage) if tiered_base_costs is not None: - return _apply_off_peak_to_base_costs(model_info, current_time, tiered_base_costs) + return _apply_off_peak_to_base_costs(model_info, current_time, *tiered_base_costs) # Get service tier aware cost keys input_cost_key: Final = get_service_tier_cost_key("input_cost_per_token", service_tier) @@ -734,13 +752,13 @@ def _get_token_base_cost( # For image generation models that don't have output_cost_per_token, # use output_cost_per_image_token as the base cost (all output tokens are image tokens) - if completion_base_cost == 0.0 or completion_base_cost is None: + if completion_base_cost == 0.0: output_image_cost: Final = get_cost_per_unit(model_info, "output_cost_per_image_token", None) if output_image_cost is not None: completion_base_cost = output_image_cost cache_creation_cost = get_cost_per_unit(model_info, cache_creation_cost_key, default_value=None) - cache_creation_cost_above_1hr = get_cost_per_unit( - model_info, "cache_creation_input_token_cost_above_1hr", default_value=None + cache_creation_keys_by_ttl: Mapping[str, str] = MappingProxyType( + {ttl: get_service_tier_cost_key(key, service_tier) for ttl, key in _TTL_PRICED_CACHE_WRITE_RATE_KEYS.items()} ) cache_read_cost = get_cost_per_unit(model_info, cache_read_cost_key, default_value=None) @@ -811,14 +829,6 @@ def _get_token_base_cost( if service_tier else f"cache_creation_input_token_cost_above_{threshold_str}_tokens" ) - cache_creation_1hr_tiered_key = ( - get_service_tier_cost_key( - f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens", - service_tier, - ) - if service_tier - else f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens" - ) cache_read_tiered_key = ( get_service_tier_cost_key( f"cache_read_input_token_cost_above_{threshold_str}_tokens", @@ -830,10 +840,19 @@ def _get_token_base_cost( cache_creation_cost = get_cost_per_unit(model_info, cache_creation_tiered_key, cache_creation_cost) - cache_creation_cost_above_1hr = get_cost_per_unit( - model_info, cache_creation_1hr_tiered_key, cache_creation_cost_above_1hr + cache_creation_keys_by_ttl = MappingProxyType( + { + ttl: _first_present_key( + model_info, + get_service_tier_cost_key( + f"{_TTL_PRICED_CACHE_WRITE_RATE_KEYS[ttl]}_above_{threshold_str}_tokens", + service_tier, + ), + key, + ) + for ttl, key in cache_creation_keys_by_ttl.items() + } ) - cache_read_cost = get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost) except (IndexError, ValueError): continue @@ -849,16 +868,26 @@ def _get_token_base_cost( resolved_cache_creation_cost: Final = ( input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost ) - - return _apply_off_peak_to_base_costs( + base_rates: Final = apply_off_peak_pricing( model_info, current_time, - ( - prompt_base_cost, - completion_base_cost, - resolved_cache_creation_cost, - cache_creation_cost_above_1hr, - resolved_cache_read_cost, + TokenRates( + input_rate=prompt_base_cost, + output_rate=completion_base_cost, + cache_read_rate=resolved_cache_read_cost, + cache_creation_rate=resolved_cache_creation_cost, + reasoning_rate=None, + ), + ) + return replace( + base_rates, + cache_creation_rates_by_ttl=MappingProxyType( + { + ttl: base_rates.cache_creation_rate + if (rate := get_cost_per_unit(model_info, key, default_value=None)) is None + else rate + for ttl, key in cache_creation_keys_by_ttl.items() + } ), ) @@ -975,25 +1004,28 @@ def resolve_image_model_info(model: str, custom_llm_provider: str, model_info: M def calculate_cache_writing_cost( cache_creation_tokens: int, cache_creation_token_details: CacheCreationTokenDetails | None, - cache_creation_cost_above_1hr: float, cache_creation_cost: float, + cache_creation_cost_by_ttl: Mapping[str, float], ) -> float: """ - Adjust cost of cache creation tokens based on the cache creation token details. + Bill cache writes per lifetime: Anthropic reports 5m and 1h counts, OpenAI reports 30m and 12h counts. + Writes no lifetime count covers, such as a lifetime a provider adds later, bill at the base write rate. """ - total_cost: float = 0.0 - if cache_creation_token_details is not None: - # get the number of 5m and 1h cache creation tokens - cache_creation_tokens_5m: Final = cache_creation_token_details.ephemeral_5m_input_tokens - cache_creation_tokens_1h: Final = cache_creation_token_details.ephemeral_1h_input_tokens - # add the number of 5m and 1h cache creation tokens to the cache creation tokens - total_cost += cache_creation_tokens_5m * cache_creation_cost if cache_creation_tokens_5m is not None else 0.0 - total_cost += ( - cache_creation_tokens_1h * cache_creation_cost_above_1hr if cache_creation_tokens_1h is not None else 0.0 - ) - else: - total_cost += cache_creation_tokens * cache_creation_cost - return total_cost + if cache_creation_token_details is None: + return cache_creation_tokens * cache_creation_cost + tokens_by_ttl: Final[Mapping[str, int | None]] = MappingProxyType( + { + "5m": cache_creation_token_details.ephemeral_5m_input_tokens, + "30m": cache_creation_token_details.ephemeral_30m_input_tokens, + "1hr": cache_creation_token_details.ephemeral_1h_input_tokens, + "12hr": cache_creation_token_details.ephemeral_12h_input_tokens, + } + ) + unattributed_tokens: Final = max(cache_creation_tokens - sum(t or 0 for t in tokens_by_ttl.values()), 0) + return unattributed_tokens * cache_creation_cost + sum( + (tokens or 0) * cache_creation_cost_by_ttl.get(ttl, cache_creation_cost) + for ttl, tokens in tokens_by_ttl.items() + ) class PromptTokensDetailsResult(TypedDict): @@ -1178,7 +1210,7 @@ def _calculate_input_cost( prompt_base_cost: float, cache_read_cost: float, cache_creation_cost: float, - cache_creation_cost_above_1hr: float, + cache_creation_cost_by_ttl: Mapping[str, float], service_tier: str | None = None, ) -> float: """ @@ -1248,8 +1280,8 @@ def _calculate_input_cost( prompt_cost += calculate_cache_writing_cost( cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"], cache_creation_token_details=prompt_tokens_details["cache_creation_token_details"], - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, cache_creation_cost=cache_creation_cost, + cache_creation_cost_by_ttl=cache_creation_cost_by_ttl, ) ### CHARACTER COST @@ -1510,34 +1542,31 @@ def generic_cost_per_token( ) billing_time: Final = current_time if current_time is not None else current_billing_time() - ( - prompt_base_cost, - completion_base_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, - cache_read_cost, - ) = _get_token_base_cost( + base_costs: Final = _get_token_base_cost( model_info=resolved_model_info, usage=usage, service_tier=service_tier, current_time=billing_time, threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) + completion_base_cost: Final = base_costs.output_rate resolved_cache_read_cost: Final = resolve_explicit_cache_read_cost_rate( model_info=resolved_model_info, usage=usage, current_time=billing_time, service_tier=service_tier, ) - effective_cache_read_cost: Final = cache_read_cost if resolved_cache_read_cost is None else resolved_cache_read_cost + effective_cache_read_cost: Final = ( + base_costs.cache_read_rate if resolved_cache_read_cost is None else resolved_cache_read_cost + ) prompt_cost = _calculate_input_cost( prompt_tokens_details=prompt_tokens_details, model_info=resolved_model_info, - prompt_base_cost=prompt_base_cost, + prompt_base_cost=base_costs.input_rate, cache_read_cost=effective_cache_read_cost, - cache_creation_cost=cache_creation_cost, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, + cache_creation_cost=base_costs.cache_creation_rate, + cache_creation_cost_by_ttl=base_costs.cache_creation_rates_by_ttl, service_tier=service_tier, ) @@ -1640,7 +1669,7 @@ class BilledTokenRates: cache_read_input_token_cost: float cache_read_input_audio_token_cost: float cache_creation_input_token_cost: float - cache_creation_input_token_cost_above_1hr: float + cache_creation_input_token_cost_by_ttl: Mapping[str, float] output_cost_per_reasoning_token: float def scaled(self, multiplier: float) -> "BilledTokenRates": @@ -1652,7 +1681,9 @@ class BilledTokenRates: cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier, cache_read_input_audio_token_cost=self.cache_read_input_audio_token_cost * multiplier, cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier, - cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier, + cache_creation_input_token_cost_by_ttl=MappingProxyType( + {ttl: rate * multiplier for ttl, rate in self.cache_creation_input_token_cost_by_ttl.items()} + ), output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier, ) @@ -1702,7 +1733,7 @@ def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRat cache_read_input_token_cost=cache_read_rate, cache_read_input_audio_token_cost=cache_read_rate, cache_creation_input_token_cost=cache_creation_rate, - cache_creation_input_token_cost_above_1hr=cache_creation_rate, + cache_creation_input_token_cost_by_ttl=MappingProxyType({}), output_cost_per_reasoning_token=output_rate, ) @@ -1717,13 +1748,7 @@ def _cost_map_billed_rates( current_time: datetime | None, ) -> BilledTokenRates: billing_time: Final = current_time if current_time is not None else current_billing_time() - ( - prompt_base_cost, - completion_base_cost, - cache_creation_cost_rate, - cache_creation_cost_above_1hr_rate, - cache_read_cost_rate, - ) = _get_token_base_cost( + base_costs: Final = _get_token_base_cost( model_info=model_info, usage=usage, service_tier=service_tier, @@ -1737,14 +1762,14 @@ def _cost_map_billed_rates( service_tier=service_tier, ) effective_cache_read_cost_rate: Final = ( - cache_read_cost_rate if provider_cache_read_cost_rate is None else provider_cache_read_cost_rate + base_costs.cache_read_rate if provider_cache_read_cost_rate is None else provider_cache_read_cost_rate ) reasoning_rate: Final = _resolve_billed_reasoning_rate( model_info=model_info, usage=usage, service_tier=service_tier, - completion_base_cost=completion_base_cost, + completion_base_cost=base_costs.output_rate, current_time=billing_time, ) audio_cache_read_rate: Final = get_cost_per_unit( @@ -1759,14 +1784,14 @@ def _cost_map_billed_rates( * get_provider_specific_speed_multiplier(model_info=model_info, usage=usage) ) return BilledTokenRates( - input_cost_per_token=prompt_base_cost, - output_cost_per_token=completion_base_cost, + input_cost_per_token=base_costs.input_rate, + output_cost_per_token=base_costs.output_rate, cache_read_input_token_cost=effective_cache_read_cost_rate, cache_read_input_audio_token_cost=( audio_cache_read_rate if audio_cache_read_rate is not None else effective_cache_read_cost_rate ), - cache_creation_input_token_cost=cache_creation_cost_rate, - cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate, + cache_creation_input_token_cost=base_costs.cache_creation_rate, + cache_creation_input_token_cost_by_ttl=base_costs.cache_creation_rates_by_ttl, output_cost_per_reasoning_token=reasoning_rate, ).scaled(multiplier) @@ -1854,8 +1879,8 @@ def get_token_type_cost_breakdown( else calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, - cache_creation_cost_above_1hr=rates.cache_creation_input_token_cost_above_1hr, cache_creation_cost=rates.cache_creation_input_token_cost, + cache_creation_cost_by_ttl=rates.cache_creation_input_token_cost_by_ttl, ) ) return TokenTypeCostBreakdown( @@ -1886,7 +1911,7 @@ def calculate_prompt_caching_savings( biller saw them rather than at the later spend write. """ model_info_with_cache_read_default: Final = apply_provider_cache_read_default(model_info, custom_llm_provider) - prompt_base_cost, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = _get_token_base_cost( + base_costs: Final = _get_token_base_cost( model_info=model_info_with_cache_read_default, usage=usage, service_tier=service_tier, @@ -1899,9 +1924,11 @@ def calculate_prompt_caching_savings( current_time=billed_at, service_tier=service_tier, ) - effective_cache_read_cost: Final = cache_read_cost if provider_cache_read_cost is None else provider_cache_read_cost - write_rate: Final = cache_creation_cost or prompt_base_cost - write_rate_1h: Final = cache_creation_cost_above_1hr or write_rate + prompt_base_cost: Final = base_costs.input_rate + effective_cache_read_cost: Final = ( + base_costs.cache_read_rate if provider_cache_read_cost is None else provider_cache_read_cost + ) + write_rate: Final = base_costs.cache_creation_rate or prompt_base_cost prompt_tokens_details: Final = parse_prompt_tokens_details(usage) cache_read_tokens: Final = max(prompt_tokens_details["cache_hit_tokens"], 0) cache_creation_tokens: Final = max(prompt_tokens_details["cache_creation_tokens"], 0) @@ -1909,7 +1936,9 @@ def calculate_prompt_caching_savings( cache_creation_details: Final = ( CacheCreationTokenDetails( ephemeral_5m_input_tokens=max(details.ephemeral_5m_input_tokens or 0, 0), + ephemeral_30m_input_tokens=max(details.ephemeral_30m_input_tokens or 0, 0), ephemeral_1h_input_tokens=max(details.ephemeral_1h_input_tokens or 0, 0), + ephemeral_12h_input_tokens=max(details.ephemeral_12h_input_tokens or 0, 0), ) if details is not None else None @@ -1918,8 +1947,10 @@ def calculate_prompt_caching_savings( write_premium: Final = calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_details, - cache_creation_cost_above_1hr=write_rate_1h - prompt_base_cost, cache_creation_cost=write_rate - prompt_base_cost, + cache_creation_cost_by_ttl=MappingProxyType( + {ttl: rate - prompt_base_cost for ttl, rate in base_costs.cache_creation_rates_by_ttl.items()} + ), ) uplift: Final = get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift( model_info, vertex_location diff --git a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py index 9688b511ea8..3ed598b5e97 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py +++ b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py @@ -5,6 +5,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never +from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.utils import ( CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, @@ -16,6 +17,8 @@ ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total" _TEXT_INPUT_RATE: Final = "input_cost_per_token" _AUDIO_INPUT_RATE: Final = "input_cost_per_audio_token" +_CACHE_READ_INPUT_RATE: Final = "cache_read_input_token_cost" +_CACHE_READ_AUDIO_INPUT_RATE: Final = "cache_read_input_audio_token_cost" _TEXT_OUTPUT_RATE: Final = "output_cost_per_token" _AUDIO_OUTPUT_RATE: Final = "output_cost_per_audio_token" _RATE_KEY_MARKERS: Final = ("cost", "pricing") @@ -33,19 +36,38 @@ def _tokens(value: object) -> int: def used_pricing_keys(usage: Usage) -> tuple[str, ...]: - prompt_audio: Final = _audio_tokens(usage.prompt_tokens_details) + prompt_details: Final = parse_prompt_tokens_details(usage) + prompt_audio: Final = prompt_details["audio_tokens"] completion_audio: Final = _audio_tokens(usage.completion_tokens_details) - prompt_text: Final = _tokens(usage.prompt_tokens) - prompt_audio + prompt_text: Final = max( + _tokens(usage.prompt_tokens) + - prompt_details["cache_hit_tokens"] + - prompt_audio + - prompt_details["image_tokens"] + - prompt_details["video_tokens"], + 0, + ) + cached_text: Final = max(prompt_details["cache_hit_tokens"] - prompt_details["cache_hit_audio_tokens"], 0) completion_text: Final = _tokens(usage.completion_tokens) - completion_audio components: Final = ( (_TEXT_INPUT_RATE, prompt_text), (_AUDIO_INPUT_RATE, prompt_audio), + (_CACHE_READ_INPUT_RATE, cached_text), + (_CACHE_READ_AUDIO_INPUT_RATE, prompt_details["cache_hit_audio_tokens"]), (_TEXT_OUTPUT_RATE, completion_text), (_AUDIO_OUTPUT_RATE, completion_audio), ) return tuple(key for key, count in components if count > 0) +def _effective_pricing_key(key: str, pricing_entry: Mapping[str, object]) -> str: + if key == _CACHE_READ_AUDIO_INPUT_RATE and pricing_entry.get(key) is None: + return _CACHE_READ_INPUT_RATE if pricing_entry.get(_CACHE_READ_INPUT_RATE) is not None else _TEXT_INPUT_RATE + if key == _CACHE_READ_INPUT_RATE and pricing_entry.get(key) is None: + return _TEXT_INPUT_RATE + return key + + def _nested_pricing(value: object) -> Mapping[str, object] | tuple[object, ...] | None: try: return _NESTED_PRICING.validate_python(value) @@ -89,7 +111,8 @@ def diagnose_zero_cost( pricing_entry: Mapping[str, object], calculation_failed: bool, ) -> StandardLoggingZeroCostDiagnostic | None: - used_keys: Final = used_pricing_keys(usage) + effective_keys: Final = tuple(_effective_pricing_key(key, pricing_entry) for key in used_pricing_keys(usage)) + used_keys: Final = tuple(key for index, key in enumerate(effective_keys) if key not in effective_keys[:index]) if not used_keys: return None missing_keys: Final = tuple(key for key in used_keys if pricing_entry.get(key) is None) diff --git a/litellm/litellm_core_utils/llm_judge.py b/litellm/litellm_core_utils/llm_judge.py index 0c2d7767455..307e000119b 100644 --- a/litellm/litellm_core_utils/llm_judge.py +++ b/litellm/litellm_core_utils/llm_judge.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: JSON_FENCE_RE: Final = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL | re.IGNORECASE) _OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) def default_router_provider() -> Router | None: @@ -40,13 +41,13 @@ def parse_json_verdict(raw: str) -> dict[str, object]: # mutable-ok: plain pars text = fenced.group(1).strip() # rebind-ok: progressively narrowed to the JSON payload parsed: object try: - parsed = json.loads(text) + parsed = _DECODED_JSON.validate_python(json.loads(text)) except json.JSONDecodeError: start: Final = text.find("{") end: Final = text.rfind("}") if start == -1 or end <= start: raise - parsed = json.loads(text[start : end + 1]) + parsed = _DECODED_JSON.validate_python(json.loads(text[start : end + 1])) if not isinstance(parsed, dict): raise ValueError("judge response is not a JSON object") return {str(k): v for k, v in parsed.items()} diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9f521ffc6cc..56c6436eae6 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -1,4 +1,6 @@ -from collections.abc import Callable +import threading +from collections.abc import Callable, Iterable, Iterator, Mapping +from types import MethodType from typing import TYPE_CHECKING, Final import litellm @@ -15,6 +17,11 @@ else: _custom_logger_compatible_callbacks_literal = str _generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {} +_callback_lists_lock: Final = threading.RLock() + + +def _instance_attributes(target: object) -> Mapping[str, object]: + return vars(target) class LoggingCallbackManager: @@ -43,16 +50,18 @@ class LoggingCallbackManager: Add a input callback to litellm.input_callback. Auto-routes async callbacks to litellm._async_input_callback. """ - if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) + with _callback_lists_lock: + if not isinstance(callback, str) and self._is_async_callable(callback): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) def add_litellm_service_callback(self, callback: CustomLogger | str | Callable[..., object]): """ Add a service callback to litellm.service_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) def add_litellm_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -60,10 +69,11 @@ class LoggingCallbackManager: Ensures no duplicates are added. """ - self._safe_add_callback_to_list( - callback=callback, - parent_list=litellm.callbacks, - ) + with _callback_lists_lock: + self._safe_add_callback_to_list( + callback=callback, + parent_list=litellm.callbacks, + ) def add_litellm_success_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -71,37 +81,41 @@ class LoggingCallbackManager: Auto-routes async callbacks to litellm._async_success_callback. Special-cases 'dynamodb' and 'openmeter' as async callbacks. """ - if ( - isinstance(callback, str) - and callback in ("dynamodb", "openmeter") - or not isinstance(callback, str) - and self._is_async_callable(callback) - ): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) + with _callback_lists_lock: + if ( + isinstance(callback, str) + and callback in ("dynamodb", "openmeter") + or not isinstance(callback, str) + and self._is_async_callable(callback) + ): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) def add_litellm_failure_callback(self, callback: CustomLogger | str | Callable[..., object]): """ Add a failure callback to `litellm.failure_callback`. Auto-routes async callbacks to litellm._async_failure_callback. """ - if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) + with _callback_lists_lock: + if not isinstance(callback, str) and self._is_async_callable(callback): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) def add_litellm_async_success_callback(self, callback: CustomLogger | Callable[..., object] | str): """ Add a success callback to litellm._async_success_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) def add_litellm_async_failure_callback(self, callback: CustomLogger | Callable[..., object] | str): """ Add a failure callback to litellm._async_failure_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) def remove_callback_from_list_by_object(self, callback_list, obj, require_self=True): """ @@ -110,13 +124,14 @@ class LoggingCallbackManager: if not isinstance(callback_list, list): # Not list -> do nothing return - if require_self: - remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] - else: - remove_list = [c for c in callback_list if c == obj] + with _callback_lists_lock: + if require_self: + remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] + else: + remove_list = [c for c in callback_list if c == obj] - for c in remove_list: - callback_list.remove(c) + for c in remove_list: + callback_list.remove(c) def remove_callbacks_by_type(self, callback_list, callback_type): """ @@ -134,10 +149,11 @@ class LoggingCallbackManager: if not isinstance(callback_list, list): return - remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)] + with _callback_lists_lock: + remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)] - for c in remove_list: - callback_list.remove(c) + for c in remove_list: + callback_list.remove(c) def _add_string_callback_to_list( self, callback: str, parent_list: list[CustomLogger | Callable[..., object] | str] @@ -329,13 +345,62 @@ class LoggingCallbackManager: key_parts: Final = [type(custom_logger).__name__] # Add only fundamental type instance variables to the key - for attr_name, attr_value in vars(custom_logger).items(): + for attr_name, attr_value in _instance_attributes(custom_logger).items(): if not attr_name.startswith("_"): # Skip private attributes if isinstance(attr_value, (str, bool, int)): key_parts.append(f"{attr_name}={attr_value}") return "-".join(key_parts) + get_custom_logger_key = _get_custom_logger_key + callback_lists_lock: Final = _callback_lists_lock + + def drop_callbacks(self, callbacks: Iterable[object], method_owners: Iterable[object] = ()) -> frozenset[str]: + dropped_ids: Final = frozenset(map(id, callbacks)) + owner_ids: Final = frozenset(map(id, method_owners)) + + def keep(callback: object) -> bool: + if id(callback) in dropped_ids: + return False + return not (isinstance(callback, MethodType) and id(callback.__self__) in owner_ids) + + with _callback_lists_lock: + vacated_keys: Final = frozenset( + self._get_custom_logger_key(callback) + for callback in litellm.callbacks + if isinstance(callback, CustomLogger) and id(callback) in dropped_ids + ) + litellm.callbacks = [callback for callback in litellm.callbacks if keep(callback)] + litellm.input_callback = [callback for callback in litellm.input_callback if keep(callback)] + litellm.success_callback = [callback for callback in litellm.success_callback if keep(callback)] + litellm.failure_callback = [callback for callback in litellm.failure_callback if keep(callback)] + litellm.service_callback = [callback for callback in litellm.service_callback if keep(callback)] + litellm._async_success_callback = [ + callback for callback in litellm._async_success_callback if keep(callback) + ] + litellm._async_failure_callback = [ + callback for callback in litellm._async_failure_callback if keep(callback) + ] + return vacated_keys + + def with_custom_logger_keys( + self, candidates: Iterable[CustomLogger], keys: frozenset[str] + ) -> tuple[CustomLogger, ...]: + if not keys: + return () + class_names: Final = frozenset(key.split("-", 1)[0] for key in keys) + same_class: Final = (candidate for candidate in candidates if type(candidate).__name__ in class_names) + return tuple(candidate for candidate in same_class if self._get_custom_logger_key(candidate) in keys) + + def without_custom_logger_keys( + self, callbacks: Iterable[str | Callable[..., object] | CustomLogger], keys: frozenset[str] + ) -> Iterator[str | Callable[..., object] | CustomLogger]: + return ( + callback + for callback in callbacks + if not (isinstance(callback, CustomLogger) and self._get_custom_logger_key(callback) in keys) + ) + def _reset_all_callbacks(self): """ Reset all callbacks to an empty list @@ -363,19 +428,21 @@ class LoggingCallbackManager: _get_all_callbacks = get_all_callbacks - def remove_callback_from_all_lists(self, obj, require_self=False) -> None: + def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None: """ Remove a callback object from every callback list it may have been promoted into, so a re-initialized callback leaves no stale instance behind. """ - for callback_list in ( - litellm.callbacks, - litellm.success_callback, - litellm.failure_callback, - litellm._async_success_callback, - litellm._async_failure_callback, - ): - self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) + with _callback_lists_lock: + for callback_list in ( + litellm.callbacks, + litellm.input_callback, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ): + self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) def get_active_additional_logging_utils_from_custom_logger( self, diff --git a/litellm/litellm_core_utils/offline_mode.py b/litellm/litellm_core_utils/offline_mode.py new file mode 100644 index 00000000000..7fdcae8825f --- /dev/null +++ b/litellm/litellm_core_utils/offline_mode.py @@ -0,0 +1,7 @@ +import os + + +def is_offline_mode() -> bool: + """True when ``LITELLM_OFFLINE_MODE=true``. A catalog, check or tokenizer LiteLLM fetches on its own, rather + than from a destination the operator configured, must check this and use its bundled copy or skip instead""" + return os.getenv("LITELLM_OFFLINE_MODE", "").strip().lower() == "true" diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 60be4d44706..59882dcbf2f 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -20,6 +20,7 @@ from pydantic import TypeAdapter import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( AllMessageValues, @@ -1527,7 +1528,7 @@ def _grammar_branches(node: Mapping[str, object]) -> int: if properties is None: return 0 optional: Final = len(_optional_property_names(node)) - return 2**optional * (len(properties) - optional + 1) + return (1 << optional) * (len(properties) - optional + 1) def _optional_property_names(node: Mapping[str, object]) -> frozenset[str]: @@ -2839,7 +2840,7 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: break try: - obj, end_idx = decoder.raw_decode(raw, idx) + decoded = decoder.raw_decode(raw, idx) except json.JSONDecodeError as e: verbose_logger.warning( "split_concatenated_json_objects: discarding unparseable tool-call " @@ -2849,13 +2850,14 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: e, ) break - if isinstance(obj, dict): - results.append(obj) + obj = as_str_mapping(_DECODED_JSON.validate_python(decoded[0])) + if obj is not None: + results.append(dict(obj)) else: # Non-dict JSON value – wrap in empty dict (Bedrock requires # toolUse.input to be an object). results.append({}) - idx = end_idx + idx = decoded[1] return results @@ -2887,14 +2889,16 @@ def salvage_concatenated_tool_arguments(raw: str) -> tuple[dict[str, object], .. if idx >= length: break try: - obj, end_idx = decoder.raw_decode(stripped, idx) + decoded = decoder.raw_decode(stripped, idx) except json.JSONDecodeError: return () - if not isinstance(obj, dict): + decoded_object = as_str_mapping(_DECODED_JSON.validate_python(decoded[0])) + if decoded_object is None: return () - idx = end_idx + idx = decoded[1] if exceeded: continue + obj = dict(decoded_object) count += 1 if not kept: kept = (obj,) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3d865bd1819..3bc8dc4e23e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -21,6 +21,7 @@ import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid from litellm.constants import REDACTED_BY_LITELLM +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_messages, find_anthropic_server_tool_result, @@ -1021,11 +1022,12 @@ def convert_to_anthropic_tool_invoke_xml(tool_calls: list) -> str: tool_function = get_attribute_or_key(tool, "function") tool_name = get_attribute_or_key(tool_function, "name") tool_arguments = get_attribute_or_key(tool_function, "arguments") - parsed_args = parse_tool_call_arguments( - tool_arguments, tool_name=tool_name, context="Anthropic XML tool invoke" + parsed_args = _DECODED_JSON.validate_python( + parse_tool_call_arguments(tool_arguments, tool_name=tool_name, context="Anthropic XML tool invoke") ) - if isinstance(parsed_args, dict): - parameters = "".join(f"<{param}>{val}\n" for param, val in parsed_args.items()) + parsed_mapping = as_str_mapping(parsed_args) + if parsed_mapping is not None: + parameters = "".join(f"<{param}>{val}\n" for param, val in parsed_mapping.items()) else: parameters = f"{parsed_args}\n" invokes += f"\n{tool_name}\n\n{parameters}\n\n" @@ -1488,6 +1490,7 @@ def convert_to_gemini_tool_call_result( message: ChatCompletionToolMessage | ChatCompletionFunctionMessage, last_message_with_tool_calls: dict | None, forward_function_call_id: bool = False, + supports_multimodal_function_responses: bool = True, ) -> VertexPartType | list[VertexPartType]: """ OpenAI message with a tool result looks like: @@ -1650,9 +1653,9 @@ def convert_to_gemini_tool_call_result( _part: Final[VertexPartType] = {"function_response": _function_response} - # For multimodal function responses, Gemini expects media parts nested - # inside functionResponse.parts instead of sibling content parts. if inline_data_list: + if not supports_multimodal_function_responses: + return [_part, *(VertexPartType(inline_data=inline_data) for inline_data in inline_data_list)] _function_response["parts"] = [{"inline_data": inline_data} for inline_data in inline_data_list] return [_part] @@ -2203,15 +2206,97 @@ def anthropic_process_openai_file_message( _EMPTY_TEXT_PLACEHOLDER: Final = "[System: Empty message content sanitised to satisfy protocol]" +def _dict_block(block: object) -> Mapping[str, object] | None: + if not isinstance(block, dict): + return None + return cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance; the converters only read dict blocks + + +def _is_empty_text_block(block: object) -> bool: + match _dict_block(block): + case {"type": "text", "text": str() as text}: + return not text.strip() + case {"type": "text"}: + return True + case _: + return False + + +def _with_placeholder_text(block: object) -> object: + fields: Final = _dict_block(block) + if fields is None or not _is_empty_text_block(fields): + return block + return {**fields, "text": _EMPTY_TEXT_PLACEHOLDER} + + +def _survives_assistant_conversion(block: object) -> bool: + """Whether ``anthropic_messages_pt`` keeps this list-content block of an assistant turn. + + Mirrors the assistant converter: a ``thinking`` / ``redacted_thinking`` block is kept + unless it is unsignable (a signed block is kept even with empty text), and + ``server_tool_use`` and ``*_tool_result`` blocks pass through. A redacted block with + blank data is dropped under ``drop_params``, which is not known here, so it never counts. + """ + fields: Final = _dict_block(block) + if fields is None: + return False + if (thinking_block := _as_thinking_block(fields)) is not None: + return not (is_unsignable_thinking_block(thinking_block) or has_blank_redacted_data(thinking_block)) + match fields: + case {"type": "server_tool_use"}: + return True + case {"type": str() as block_type}: + return block_type.endswith("_tool_result") + case _: + return False + + +def _tool_call_survives_conversion(tool_call: object) -> bool: + """Whether ``convert_to_anthropic_tool_invoke`` can emit a block for this tool call. + + Non-function calls are skipped there, and a tool search call is dropped when its + result is not replayed, which depends on the rest of the conversation, so neither counts. + """ + if message_field(tool_call, "type") != "function": + return False + match message_field(tool_call, "id"), message_field(message_field(tool_call, "function"), "name"): + case (str() as tool_id, str() as tool_name): + return not is_anthropic_tool_search_use(tool_id, tool_name) + case _: + return True + + +def _carries_non_text_content(message: AllMessageValues) -> bool: + if any(_tool_call_survives_conversion(tool_call) for tool_call in parts_of(message_field(message, "tool_calls"))): + return True + content: Final = message_field(message, "content") + blocks: Final = parts_of(content) + return isinstance(content, list) and any(_survives_assistant_conversion(block) for block in blocks) + + +def _with_content(message: AllMessageValues, content: object, action: str) -> AllMessageValues: + sanitized: Final = cast(AllMessageValues, {**message, "content": content}) + verbose_logger.debug("_sanitize_empty_text_content: %s in %s message", action, message_field(message, "role")) + return sanitized + + def _sanitize_empty_text_content( message: AllMessageValues, ) -> AllMessageValues: """ Case C: Sanitize empty text content - - Replace empty or whitespace-only text content with a placeholder message. - - Handles both string content and list-of-blocks content (rewriting only - the empty text blocks in place; non-text blocks like images are left - untouched). + + Anthropic rejects empty text blocks, so empty or whitespace-only text must never be + forwarded verbatim. What replaces it depends on what else the turn carries. + + Assistant turns that already carry substantive content (``tool_calls``, or non-text + blocks such as ``server_tool_use``) get the empty text dropped. Anthropic accepts an + assistant turn whose content is just ``tool_use``, so there is nothing to substitute, and + substituting anyway leaks a protocol diagnostic into the conversation as though the + assistant had said it, which the model then echoes on later turns. + + Every other empty turn would otherwise end up with no content at all, which Anthropic + also rejects, so the placeholder is still substituted there. Returns: The message with sanitized content if needed, otherwise the original message @@ -2219,42 +2304,26 @@ def _sanitize_empty_text_content( if message.get("role") not in ["user", "assistant"]: return message - content: Final = message.get("content") + content: Final = message_field(message, "content") + drop_empty_text: Final = message_field(message, "role") == "assistant" and _carries_non_text_content(message) if isinstance(content, str): - if not content or not content.strip(): - message = cast(AllMessageValues, dict(message)) # Make a copy - message["content"] = _EMPTY_TEXT_PLACEHOLDER - verbose_logger.debug( - "_sanitize_empty_text_content: Replaced empty text content in %s message", message.get("role") - ) + if content.strip(): + return message + if drop_empty_text: + return _with_content(message, None, "Dropped empty text content") + return _with_content(message, _EMPTY_TEXT_PLACEHOLDER, "Replaced empty text content") + + blocks: Final = parts_of(content) + if not isinstance(content, list) or not any(_is_empty_text_block(block) for block in blocks): return message - if isinstance(content, list): - # Walk the blocks and rewrite any empty text blocks. We rewrite (rather - # than drop) so callers don't end up with an entirely empty content - # list, which Anthropic also rejects. - new_blocks: Final[list[Any]] = [] - rewrote_any = False - for block in content: - if isinstance(block, dict) and block.get("type") == "text": - text = block.get("text") - if not isinstance(text, str) or not text or not text.strip(): - new_block = dict(block) - new_block["text"] = _EMPTY_TEXT_PLACEHOLDER - new_blocks.append(new_block) - rewrote_any = True - continue - new_blocks.append(block) + if drop_empty_text: + kept: Final = [block for block in blocks if not _is_empty_text_block(block)] + return _with_content(message, kept, "Dropped empty text block(s)") - if rewrote_any: - message = cast(AllMessageValues, dict(message)) # Make a copy - message["content"] = new_blocks - verbose_logger.debug( - "_sanitize_empty_text_content: Replaced empty text block(s) in %s message", message.get("role") - ) - - return message + rewritten: Final = [_with_placeholder_text(block) for block in blocks] + return _with_content(message, rewritten, "Replaced empty text block(s)") def _add_missing_tool_results( @@ -2342,7 +2411,7 @@ def _add_missing_tool_results( else: tool_name = getattr(function, "name", "unknown_tool") else: - function = getattr(tool_call, "function", None) + function = attribute_or_none(tool_call, "function") if function: tool_name = getattr(function, "name", "unknown_tool") break @@ -3184,10 +3253,14 @@ def anthropic_messages_pt( def extract_between_tags(tag: str, string: str, strip: bool = False) -> list[str]: ext_list = re.findall(f"<{tag}>(.+?)", string, re.DOTALL) if strip: - ext_list = [e.strip() for e in ext_list] + ext_list = list(_stripped(ext_list)) return ext_list +def _stripped(values: Iterable[str]) -> Iterator[str]: + return (value.strip() for value in values) + + def contains_tag(tag: str, string: str) -> bool: return bool(re.search(f"<{tag}>(.+?)", string, re.DOTALL)) @@ -4061,13 +4134,15 @@ def _convert_to_bedrock_tool_call_invoke( arguments_dict = {} else: try: - arguments_dict = json.loads(arguments) + decoded_arguments = as_str_mapping(_DECODED_JSON.validate_python(json.loads(arguments))) # Ensure arguments_dict is always a dict # (Bedrock requires toolUse.input to be an object). # Some providers return arguments: '""' which # json.loads decodes to a bare string. - if not isinstance(arguments_dict, dict): + if decoded_arguments is None: arguments_dict = {} + else: + arguments_dict = dict(decoded_arguments) except json.JSONDecodeError: # The model may return multiple JSON objects # concatenated in a single arguments string, e.g. @@ -6028,7 +6103,11 @@ def _parse_tool_call_arguments(raw: object, tool_name: str | None, context: str) return ({},) normalized_raw: Final = "{}" if raw == REDACTED_BY_LITELLM else raw try: - parsed: Final = parse_tool_call_arguments(normalized_raw, tool_name=tool_name, context=context) + parsed: Final = as_str_mapping( + _DECODED_JSON.validate_python( + parse_tool_call_arguments(normalized_raw, tool_name=tool_name, context=context) + ) + ) except ValueError as e: salvaged: Final = salvage_concatenated_tool_arguments(normalized_raw) if salvaged: @@ -6041,7 +6120,7 @@ def _parse_tool_call_arguments(raw: object, tool_name: str | None, context: str) return salvaged verbose_logger.warning("Failed to parse tool call arguments: %s", e) return ({},) - return (parsed,) if isinstance(parsed, dict) else ({},) + return (dict(parsed),) if parsed is not None else ({},) def _choice_tool_calls(choice: object) -> tuple[object, ...]: diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index ecf419695d0..eb3400aee8b 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -482,7 +482,7 @@ def _runs_by_anchor( ) -> Mapping[int, tuple[tuple[int, ...], ...]]: anchored: Final = tuple(sorted(((anchor, run) for run, anchor in anchors if anchor is not None), key=itemgetter(0))) return MappingProxyType( - {anchor: tuple(run for _, run in group) for anchor, group in groupby(anchored, key=itemgetter(0))} + {anchor: tuple(run for _, run in group) for anchor, group in groupby(anchored, key=lambda entry: entry[0])} ) diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index 02c4cd69159..e022839c5c7 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -44,10 +44,10 @@ def _get_value(value: object, key: str) -> object | None: return getattr(value, key, None) -def _get_provider_affinity_header_name(litellm_params: object | None) -> str | None: +def _get_provider_affinity_header_name(litellm_params: object | None, default_header: str | None = None) -> str | None: header: Final = _get_value(litellm_params, "provider_affinity_header") if litellm_params is not None else None if header is None: - return None + return default_header if not isinstance(header, str): raise TypeError("provider_affinity_header must be a string") return validate_provider_affinity_header_name(header) @@ -84,9 +84,9 @@ def get_stable_session_id(litellm_params: object | None) -> str | None: def add_provider_affinity_header( - headers: Mapping[str, object], litellm_params: object | None + headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers - header_name: Final = _get_provider_affinity_header_name(litellm_params) + header_name: Final = _get_provider_affinity_header_name(litellm_params, default_header) if header_name is None or any(key.lower() == header_name.lower() for key in headers): return dict(headers) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 0168cdad7c0..7f8a38cc8ec 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.vertex_ai.common_utils import ( redact_vertex_ai_metadata_from_litellm_params, redact_vertex_ai_metadata_from_logged_object, @@ -41,6 +42,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObject = Any +MESSAGE_REDACTION_ENABLE_HEADERS: Final = ("litellm-enable-message-redaction", "x-litellm-enable-message-redaction") + def _opted_out_of_message_logging(custom_logger: CustomLogger) -> bool: return hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True @@ -195,33 +198,45 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ return { **fields, **({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}), - **{key: None for key in ("thinking_blocks", "reasoningContentBlocks", "citations") if key in fields}, + **{ + key: None + for key in ( + "thinking_blocks", + "reasoningContentBlocks", + "citations", + "compaction_blocks", + "compaction_delta", + ) + if key in fields + }, } +def _redact_field_when_present(item: object, field: str) -> None: + if getattr(item, field, None) is not None: + setattr(item, field, REDACTED_BY_LITELLM) + + def _redact_responses_api_output(output_items): """Helper to redact ResponsesAPIResponse output items.""" for output_item in output_items: + output_type = stream_item_field(output_item, "type") if getattr(output_item, "text", None) is not None: output_item.text = REDACTED_BY_LITELLM - if hasattr(output_item, "content") and isinstance(output_item.content, list): - for content_part in output_item.content: - if getattr(content_part, "text", None) is not None: - content_part.text = REDACTED_BY_LITELLM - if getattr(content_part, "refusal", None) is not None: - content_part.refusal = REDACTED_BY_LITELLM + for content_part in stream_item_items(output_item, "content"): + _redact_field_when_present(content_part, "text") + _redact_field_when_present(content_part, "refusal") - # Redact reasoning items in output array - if hasattr(output_item, "type") and output_item.type == "reasoning": - if hasattr(output_item, "summary") and isinstance(output_item.summary, list): - for summary_item in output_item.summary: - if getattr(summary_item, "text", None) is not None: - summary_item.text = REDACTED_BY_LITELLM + if output_type == "reasoning": + for summary_item in stream_item_items(output_item, "summary"): + _redact_field_when_present(summary_item, "text") - if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"): + if output_type == "compaction" and hasattr(output_item, "encrypted_content"): + output_item.encrypted_content = REDACTED_BY_LITELLM + if output_type == "function_call" and hasattr(output_item, "arguments"): output_item.arguments = REDACTED_BY_LITELLM - if hasattr(output_item, "type") and output_item.type == "custom_tool_call" and hasattr(output_item, "input"): + if output_type == "custom_tool_call" and hasattr(output_item, "input"): output_item.input = REDACTED_BY_LITELLM @@ -262,6 +277,11 @@ def _redacted_responses_api_output_item(item: object, redacted_str: str) -> obje if item_mapping.get("type") == "reasoning" and isinstance(summary, list) else {} ), + **( + {"encrypted_content": redacted_str} + if item_mapping.get("type") == "compaction" and "encrypted_content" in item_mapping + else {} + ), **( {"arguments": redacted_str} if item_mapping.get("type") == "function_call" and "arguments" in item_mapping @@ -556,13 +576,8 @@ def should_redact_message_logging(model_call_details: dict) -> bool: # User explicitly disabled redaction via header return False - possible_enable_headers: Final = [ - "litellm-enable-message-redaction", # old header. maintain backwards compatibility - "x-litellm-enable-message-redaction", # new header - ] - is_redaction_enabled_via_header = False - for header in possible_enable_headers: + for header in MESSAGE_REDACTION_ENABLE_HEADERS: if bool(request_headers.get(header, False)): is_redaction_enabled_via_header = True break diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 52758e617d8..8c402fe9b44 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping, Sequence +from collections.abc import ItemsView, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass, field from typing import Any, Final @@ -30,6 +30,10 @@ _DEFAULT_SENSITIVE_PATTERNS: Final = frozenset( ) +def _object_items(data: Mapping[str, object]) -> ItemsView[str, object]: + return data.items() + + class SensitiveDataMasker: def __init__( self, @@ -127,7 +131,7 @@ class SensitiveDataMasker: return data masked_data: Final[dict[str, object]] = {} - for k, v in data.items(): + for k, v in _object_items(data): try: key_is_sensitive = self.is_sensitive_key(k, excluded_keys) if isinstance(v, Mapping): diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 39f675099d7..68cdc204af4 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict @@ -425,7 +426,7 @@ class ChunkProcessor: if hasattr(chunk["choices"][0], "finish_reason"): chunk_finish_reason = chunk["choices"][0].finish_reason elif "finish_reason" in chunk["choices"][0]: - chunk_finish_reason = chunk["choices"][0]["finish_reason"] + chunk_finish_reason = _CHOICE_FIELD_VALUE.validate_python(chunk["choices"][0]["finish_reason"]) if chunk_finish_reason is not None: finish_reason = chunk_finish_reason @@ -1226,6 +1227,7 @@ def concatenate_base64_list(base64_strings: Sequence[str]) -> str: _CHOICE_DICTS: Final = TypeAdapter(list[dict[str, object]]) +_CHOICE_FIELD_VALUE: Final = TypeAdapter(object) def _text_completion_chunk_from_mapping(chunk: Mapping[str, object]) -> TextCompletionResponse: @@ -1351,14 +1353,27 @@ def stream_chunk_builder_text_completion( _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value}) -def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool: +def _logged_provider(logging_obj: Optional["Logging"]) -> object: if logging_obj is None: - return False + return None provider: Final[object] = logging_obj.model_call_details.get("custom_llm_provider") - return provider in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS + return provider -def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None: +def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool: + return _logged_provider(logging_obj) in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS + + +def _stream_builder_response_cost( + response: ModelResponse, logging_obj: Optional["Logging"], chunks: Sequence[object] +) -> float | None: + if _logged_provider(logging_obj) == LlmProviders.OPENCODE_GO.value: + # usage.opencode only survives on the chunks, the assembled Usage keeps declared fields + reported: Final = ( + provider_reported_usage_cost(chunk_usage(chunk), LlmProviders.OPENCODE_GO.value) + for chunk in reversed(chunks) + ) + return next((cost for cost in reported if cost is not None), None) usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None) if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj): return float(usage_cost) @@ -1396,8 +1411,10 @@ def _set_hidden_response_cost(response: ModelResponse, response_cost: float) -> hidden_params["response_cost"] = response_cost -def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> None: - response_cost: Final = _stream_builder_response_cost(response, logging_obj) +def _set_stream_builder_response_cost( + response: ModelResponse, logging_obj: Optional["Logging"], chunks: Sequence[object] +) -> None: + response_cost: Final = _stream_builder_response_cost(response, logging_obj, chunks) if response_cost is None: return _set_hidden_response_cost(response, response_cost) @@ -1549,7 +1566,11 @@ def stream_chunk_builder( break _stamp_streaming_usage_cost(usage, response, logging_obj) - _set_stream_builder_response_cost(response, logging_obj) + _set_stream_builder_response_cost( + response, + logging_obj, + chunks, # pyright: ignore[reportUnknownArgumentType] # legacy chunk list is untyped + ) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response @@ -1733,7 +1754,11 @@ def stream_chunk_builder( break _stamp_streaming_usage_cost(usage, response, logging_obj) - _set_stream_builder_response_cost(response, logging_obj) + _set_stream_builder_response_cost( + response, + logging_obj, + chunks, # pyright: ignore[reportUnknownArgumentType] # legacy chunk list is untyped + ) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ec426d6da5..2df470edc59 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -10,7 +10,7 @@ import traceback from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, NoReturn, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx @@ -24,6 +24,7 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, get_hidden_params, @@ -37,7 +38,6 @@ from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.llms.openai import OpenAIChatCompletionChunk from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CacheCreationTokenDetails, @@ -60,6 +60,9 @@ 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 +if TYPE_CHECKING: + from .litellm_logging import BilledRoundsSpend + # Constants for special delta attribute names AUDIO_ATTRIBUTE: Final = "audio" IMAGE_ATTRIBUTE: Final = "images" @@ -70,7 +73,9 @@ _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) _CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None -_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value}) +_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset( + {LlmProviders.OPENROUTER.value, LlmProviders.OPENCODE_GO.value} +) def _next_sync_or_exhausted(it: Iterator[object]) -> object: @@ -108,14 +113,15 @@ def _prepend_upstream_choice_extras( choice.__pydantic_fields_set__.update(upstream_extras) -def _holds_only_scalars(model: BaseModel) -> bool: - fields: Final[Iterator[tuple[str, object]]] = iter(model) +def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool: return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) -def _delta_from_upstream(upstream_delta: BaseModel | None) -> Delta | None: +def _delta_from_upstream(upstream_delta: "BaseModel | _UpstreamDeltaLike | None") -> Delta | None: if upstream_delta is None: return Delta() + if not isinstance(upstream_delta, BaseModel): + return None try: return Delta(**dict(iter(upstream_delta))) except Exception: @@ -186,6 +192,33 @@ class _DeltaDumpDict(TypedDict): tool_calls: NotRequired[Sequence[Mapping[str, object]]] +class _UpstreamFunctionCallLike(Protocol): + arguments: str | None + + +class _UpstreamToolCallLike(Protocol): + function: _UpstreamFunctionCallLike + + +class _UpstreamDeltaLike(Protocol): + function_call: _UpstreamFunctionCallLike | None + tool_calls: Sequence[_UpstreamToolCallLike] | None + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + def model_dump(self) -> _DeltaDumpDict: ... + + +class _UpstreamChoiceLike(Protocol): + delta: _UpstreamDeltaLike | None + + +class _UpstreamChunkLike(Protocol): + id: str + choices: Sequence[_UpstreamChoiceLike] + system_fingerprint: str | None + + class _TextCompletionChoiceLike(Protocol): text: str finish_reason: str | None @@ -259,6 +292,8 @@ def _provider_hidden_params( class CustomStreamWrapper: + _earlier_rounds_spend: "BilledRoundsSpend | None" = None + @property def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation return self._hidden_params @@ -342,6 +377,10 @@ class CustomStreamWrapper: self.messages = getattr(logging_obj, "messages", None) self.sent_stream_usage = False self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False + self.caller_requested_stream_usage: Final = ( + self.check_send_stream_usage(getattr(logging_obj, "stream_options", None)) + and getattr(logging_obj, "stream_usage_injected", False) is not True + ) self.tool_call = False # keep track of the returned chunks - used for calculating the input/output tokens for stream options self.chunks: list[ModelResponseStream] = [] @@ -911,7 +950,7 @@ class CustomStreamWrapper: def copy_model_response_level_provider_specific_fields( self, - original_chunk: ModelResponseStream | OpenAIChatCompletionChunk, + original_chunk: object, model_response: ModelResponseStream, ) -> ModelResponseStream: """ @@ -940,8 +979,8 @@ class CustomStreamWrapper: or ("function_call" in completion_obj and completion_obj["function_call"] is not None) or ( "tool_calls" in model_response.choices[0].delta - and model_response.choices[0].delta["tool_calls"] is not None - and len(model_response.choices[0].delta["tool_calls"]) > 0 + and model_response.choices[0].delta.tool_calls is not None + and len(model_response.choices[0].delta.tool_calls) > 0 ) or ( "function_call" in model_response.choices[0].delta @@ -1106,7 +1145,8 @@ class CustomStreamWrapper: ## check if openai/azure chunk original_chunk: Final = response_obj.get("original_chunk", None) if original_chunk: - if len(original_chunk.choices) > 0: + upstream_chunk: Final[_UpstreamChunkLike] = original_chunk + if len(upstream_chunk.choices) > 0: original_chunk_for_reuse: Final[object] = cast( # cast-ok: runtime type-check before access object, original_chunk, @@ -1142,7 +1182,7 @@ class CustomStreamWrapper: ) else: choices: Final = [] - for choice in original_chunk.choices: + for choice in upstream_chunk.choices: try: if isinstance(choice, BaseModel): choice_json = choice.model_dump() @@ -1155,11 +1195,11 @@ class CustomStreamWrapper: setattr(model_response, "choices", choices) else: return - model_response.system_fingerprint = original_chunk.system_fingerprint + model_response.system_fingerprint = upstream_chunk.system_fingerprint setattr( model_response, "citations", - stored_attribute(original_chunk, "citations"), + stored_attribute(upstream_chunk, "citations"), ) preserve_upstream_non_openai_attributes( model_response=model_response, @@ -1602,7 +1642,9 @@ class CustomStreamWrapper: model_response.model = self.model ## FUNCTION CALL PARSING - original_chunk: Final = response_obj.get("original_chunk") if response_obj is not None else None + original_chunk: Final[_UpstreamChunkLike | None] = ( + response_obj.get("original_chunk") if response_obj is not None else None + ) if ( original_chunk is not None ): # function / tool calling branch - only set for openai/azure compatible endpoints @@ -1614,31 +1656,19 @@ class CustomStreamWrapper: original_chunk, model_response ) if original_chunk.choices and len(original_chunk.choices) > 0: - delta = original_chunk.choices[0].delta + delta: Final = original_chunk.choices[0].delta if delta is not None and (delta.function_call is not None or delta.tool_calls is not None): try: model_response.system_fingerprint = original_chunk.system_fingerprint ## AZURE - check if arguments is not None - if original_chunk.choices[0].delta.function_call is not None: - if ( - getattr( - original_chunk.choices[0].delta.function_call, - "arguments", - ) - is None - ): - original_chunk.choices[0].delta.function_call.arguments = "" - elif original_chunk.choices[0].delta.tool_calls is not None: - if isinstance(original_chunk.choices[0].delta.tool_calls, list): - for t in original_chunk.choices[0].delta.tool_calls: - if hasattr(t, "functions") and hasattr(t.functions, "arguments"): - if ( - getattr( - t.function, - "arguments", - ) - is None - ): + if delta.function_call is not None: + if delta.function_call.arguments is None: + delta.function_call.arguments = "" + elif delta.tool_calls is not None: + if isinstance(delta.tool_calls, list): + for t in delta.tool_calls: + if hasattr(t, "functions") and hasattr(getattr(t, "functions"), "arguments"): + if t.function.arguments is None: t.function.arguments = "" _json_delta: Final[_DeltaDumpDict] = delta.model_dump() if "role" not in _json_delta or _json_delta["role"] is None: @@ -1877,7 +1907,7 @@ class CustomStreamWrapper: else: asyncio.run(self.logging_obj.async_success_handler(processed_chunk, None, None, cache_hit)) ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler - litellm_params: Final = self.logging_obj.model_call_details.get("litellm_params", {}) + litellm_params: Final[Mapping[str, object]] = self.logging_obj.model_call_details.get("litellm_params", {}) if self.logging_obj.is_sync_litellm_request(litellm_params): self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) @@ -1926,11 +1956,17 @@ class CustomStreamWrapper: def _propagate_usage_cost_to_hidden_params( response: "ModelResponse", custom_llm_provider: str | None, + chunks: Sequence[object] = (), ) -> None: + """The builder's usage only keeps declared fields, so provider blocks like usage.opencode are read off the chunks.""" if custom_llm_provider not in _USAGE_COST_HEADER_PROVIDERS: return - _usage: Final[Usage | None] = getattr(response, "usage", None) - _cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None)) + usages: Final = ( + getattr(response, "usage", None), + *(chunk_usage(chunk) for chunk in reversed(chunks)), + ) + costs: Final = (provider_reported_usage_cost(usage, custom_llm_provider) for usage in usages) + _cost: Final = next((cost for cost in costs if cost is not None), None) if _cost is not None: if "additional_headers" not in response.hidden_params: response.hidden_params["additional_headers"] = {} @@ -2070,9 +2106,11 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: complete_streaming_response: Final = self._build_complete_streaming_response() - response = self.model_response_creator() + response = self._new_usage_chunk() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) + self._propagate_usage_cost_to_hidden_params( + complete_streaming_response, self.custom_llm_provider, self.chunks + ) setattr( response, @@ -2222,27 +2260,10 @@ class CustomStreamWrapper: ) if _has_usage: - # Store a copy ONLY when usage stripping below will mutate - # the chunk. For non-usage chunks (vast majority), store - # directly to avoid expensive model_copy() per chunk. - self.chunks.append(processed_chunk.model_copy()) - - # Strip usage from the outgoing chunk so it's not sent twice - # (once in the chunk, once in _hidden_params). - obj_dict = processed_chunk.model_dump() - if "usage" in obj_dict: - del obj_dict["usage"] - processed_chunk = self.model_response_creator( - chunk=obj_dict, - hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage - Mapping[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR) - ), - ) - is_empty = is_model_response_stream_empty( - model_response=cast(ModelResponseStream, processed_chunk) - ) - if is_empty: + usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) + if usage_free_chunk is None: continue + processed_chunk = usage_free_chunk else: # No usage data — safe to store directly without copying self.chunks.append(processed_chunk) @@ -2273,25 +2294,25 @@ class CustomStreamWrapper: if processed_chunk is None: continue - choice = processed_chunk.choices[0] - if isinstance(choice, StreamingChoices): - self._accumulate_streamed_delta(choice.delta) + if processed_chunk.choices: + choice = processed_chunk.choices[0] + if isinstance(choice, StreamingChoices): + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) - # RETURN RESULT - self.chunks.append(processed_chunk) - return processed_chunk + if not self.send_stream_usage or getattr(processed_chunk, "usage", None) is None: + self.chunks.append(processed_chunk) + return processed_chunk + usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) + if usage_free_chunk is None: + continue + return usage_free_chunk except (StopAsyncIteration, StopIteration): return await self._finalize_completed_stream(cache_hit=cache_hit) except httpx.TimeoutException as e: # if httpx read timeout error occues traceback_exception = traceback.format_exc() ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}" - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - self._failure_logging = asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) - ) + self._dispatch_stream_failure(e, traceback_exception) await self.aclose() self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: @@ -2306,9 +2327,11 @@ class CustomStreamWrapper: async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": if self.sent_last_chunk is True: complete_streaming_response: Final = await self._async_build_complete_streaming_response() - response: Final = self.model_response_creator() + response: Final = self._new_usage_chunk() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) + self._propagate_usage_cost_to_hidden_params( + complete_streaming_response, self.custom_llm_provider, self.chunks + ) setattr( response, @@ -2392,29 +2415,36 @@ class CustomStreamWrapper: return processed_chunk def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn: - traceback_exception: Final = traceback.format_exc() - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - self._failure_logging = asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) - ) + self._dispatch_stream_failure(e, traceback.format_exc()) self._handle_stream_fallback_error(e) - async def wait_for_failure_logging(self, timeout: float) -> None: + def _dispatch_stream_failure(self, e: Exception, traceback_exception: str) -> None: + self._record_partial_usage_for_failure() + self._failure_logging = asyncio.create_task( + self.logging_obj.dispatch_attempt_failure_handlers(e, traceback_exception) + ) + + @property + def has_dispatched_failure_logging(self) -> bool: + return self._failure_logging is not None + + async def wait_for_failure_logging(self, timeout: float | None) -> None: """Let this stream's failure log read its own deployment before a fallback reuses the logging object.""" if self._failure_logging is None: return await asyncio.wait((self._failure_logging,), timeout=timeout) + def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: + self._earlier_rounds_spend = earlier + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks already delivered. Recover that partial usage from the chunks seen so far and stash it, with its cost, on the logging object so the failure - handler records the real partial spend instead of zero. A request that - later recovers via a router fallback overwrites this with the combined - success log on the same request id, so this never double counts. + handler records the real partial spend instead of zero. The router + moves it into terminal failure usage before opening the next attempt, + so a recovered attempt logs its own usage. """ if self.logging_obj is None or not self.chunks: return @@ -2433,10 +2463,12 @@ class CustomStreamWrapper: if self.model: partial_response.model = self.model backfill_missing_cache_usage_fields(usage) - self.logging_obj.model_call_details["combined_usage_object"] = usage - self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 - ) + partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 + if self._earlier_rounds_spend is None: + self.logging_obj.model_call_details["combined_usage_object"] = usage + self.logging_obj.model_call_details["response_cost"] = partial_cost + return + self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, usage, partial_cost) except Exception as recover_error: verbose_logger.debug( "could not recover partial usage for interrupted stream: %s", @@ -2463,6 +2495,22 @@ class CustomStreamWrapper: return False return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) + def _new_usage_chunk(self) -> ModelResponseStream: + usage_chunk: Final = self.model_response_creator() + if self.caller_requested_stream_usage: + usage_chunk.choices = [] + return usage_chunk + + def _record_and_strip_usage_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: + self.chunks.append(chunk.model_copy()) + usage_free_chunk: Final = self.model_response_creator( + chunk={key: value for key, value in chunk.model_dump().items() if key != "usage"}, + hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage + Mapping[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR) + ), + ) + return None if is_model_response_stream_empty(model_response=usage_free_chunk) else usage_free_chunk + def _accumulate_streamed_delta(self, delta: object) -> None: # Shared by the sync, async, and non-aiohttp iteration sites so answer # text and the disqualifying-content latch stay in step across all three. diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 8aded69aed2..90e79d4d853 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -44,7 +44,7 @@ from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEnc if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace from litellm.litellm_core_utils.url_utils import safe_get -from litellm.llms.custom_httpx.http_handler import get_httpx_client +from litellm.llms.custom_httpx.http_handler import get_httpx_client, header_value from litellm.rust_bridge.tokenizer import get_encoding from litellm.types.llms.anthropic import ( AnthropicContentParamSource, @@ -247,7 +247,7 @@ def get_image_dimensions( client: Final = get_httpx_client() response: Final[httpx.Response] = safe_get(client, data) max_bytes: Final = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) - content_length: Final[str | None] = response.headers.get("Content-Length") + content_length: Final[str | None] = header_value(response.headers, "Content-Length") if content_length is not None and int(content_length) > max_bytes: pass # skip download; img_data stays None else: diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 43e16599bf6..0db6a41b503 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -22,11 +22,11 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): import asyncio import socket from ipaddress import ip_address, ip_network -from typing import Any, Final, Protocol -from urllib.parse import quote, urlparse, urlunparse +from typing import Final, Protocol +from urllib.parse import quote, urlparse, urlsplit, urlunparse import httpx -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, TypedDict, Unpack import litellm @@ -50,6 +50,7 @@ class _UrlFetcher(Protocol): *, headers: dict[str, str] | None = None, follow_redirects: bool = False, + timeout: float | httpx.Timeout | None = None, ) -> httpx.Response: ... @@ -62,6 +63,8 @@ class _AsyncUrlFetcher(Protocol): *, headers: dict[str, str] | None = None, follow_redirects: bool = False, + timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, ) -> httpx.Response: ... @@ -73,8 +76,14 @@ class _AsyncFetcherView(TypedDict): fetcher: ReadOnly[_AsyncUrlFetcher] -class _CallerHeadersView(TypedDict): +class _FetchOptions(TypedDict, total=False): headers: ReadOnly[dict[str, str]] + follow_redirects: ReadOnly[bool] + timeout: ReadOnly[float | httpx.Timeout | None] + + +class _AsyncFetchOptions(_FetchOptions, total=False): + max_response_bytes: ReadOnly[int | None] # Cloud-internal IPs that ``ip.is_global`` can report as public. Everything @@ -129,6 +138,28 @@ def encode_url_path_segments(value: object, *, field_name: str = "path") -> str: return "/".join(encode_url_path_segment(segment, field_name=field_name) for segment in value_str.split("/")) +def _strip_trailing_slashes_keeping_the_scheme(api_base: str) -> str: + stripped: Final = api_base.rstrip("/") + parts: Final = urlsplit(stripped) + if parts.scheme and not (parts.netloc or parts.path or parts.query or parts.fragment): + return api_base + if not parts.query or not parts.path.endswith("/"): + return stripped + url_before_query, _, query = stripped.partition("?") + return f"{url_before_query.rstrip('/')}?{query}" + + +def join_api_base(api_base: str, route: str) -> str: + return f"{_strip_trailing_slashes_keeping_the_scheme(api_base)}/{route.lstrip('/')}" + + +def ensure_api_route(api_base: str, route: str) -> str: + base: Final = _strip_trailing_slashes_keeping_the_scheme(api_base) + if urlsplit(base).path.endswith(f"/{route.lstrip('/')}"): + return base + return join_api_base(base, route) + + def _is_blocked_ip(addr: str) -> bool: """Return True for any IP not safe to reach from a user-supplied URL. @@ -418,7 +449,7 @@ def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: return str(httpx.URL(request_url).join(location)) -def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: +def safe_get(client: _UrlFetcher, url: str, **kwargs: Unpack[_FetchOptions]) -> httpx.Response: """ Fetch a user-supplied URL with SSRF protection on every redirect hop. @@ -438,21 +469,23 @@ def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: The final httpx.Response. """ if not getattr(litellm, "user_url_validation", True): - kwargs.setdefault("follow_redirects", True) - unvalidated: Final[_ResponseView] = {"response": client.get(url, **kwargs)} + unvalidated_options: Final[_FetchOptions] = { + **kwargs, + "follow_redirects": kwargs.get("follow_redirects", True), + } + unvalidated: Final[_ResponseView] = {"response": client.get(url, **unvalidated_options)} return unvalidated["response"] fetcher_view: Final[_FetcherView] = {"fetcher": client} fetcher: Final = fetcher_view["fetcher"] - kwargs.pop("follow_redirects", None) - headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} + caller_headers: Final = kwargs.get("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = fetcher.get( - validated_url, - headers={**headers_view["headers"], "Host": original_host}, - follow_redirects=False, + hop_options: _FetchOptions = { **kwargs, - ) + "headers": {**caller_headers, "Host": original_host}, + "follow_redirects": False, + } + response = fetcher.get(validated_url, **hop_options) if not response.is_redirect: return response # Resolve the next hop against the ORIGINAL (pre-rewrite) URL so @@ -461,24 +494,26 @@ def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: raise SSRFError("Too many redirects") -async def async_safe_get(client: _AsyncUrlFetcher, url: str, **kwargs: Any) -> httpx.Response: +async def async_safe_get(client: _AsyncUrlFetcher, url: str, **kwargs: Unpack[_AsyncFetchOptions]) -> httpx.Response: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): - kwargs.setdefault("follow_redirects", True) - unvalidated: Final[_ResponseView] = {"response": await client.get(url, **kwargs)} + unvalidated_options: Final[_AsyncFetchOptions] = { + **kwargs, + "follow_redirects": kwargs.get("follow_redirects", True), + } + unvalidated: Final[_ResponseView] = {"response": await client.get(url, **unvalidated_options)} return unvalidated["response"] fetcher_view: Final[_AsyncFetcherView] = {"fetcher": client} fetcher: Final = fetcher_view["fetcher"] - kwargs.pop("follow_redirects", None) - headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} + caller_headers: Final = kwargs.get("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = await asyncio.to_thread(validate_url, url) - response = await fetcher.get( - validated_url, - headers={**headers_view["headers"], "Host": original_host}, - follow_redirects=False, + hop_options: _AsyncFetchOptions = { **kwargs, - ) + "headers": {**caller_headers, "Host": original_host}, + "follow_redirects": False, + } + response = await fetcher.get(validated_url, **hop_options) if not response.is_redirect: return response # Resolve the next hop against the ORIGINAL (pre-rewrite) URL so diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 7452a4a20ef..91f33635658 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -3,6 +3,7 @@ import os from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.types.utils import CallTypes from . import * @@ -128,7 +129,7 @@ def discover_guardrail_translation_mappings() -> dict[CallTypes, type["BaseTrans # Check for guardrail_translation_mappings dictionary if hasattr(module, "guardrail_translation_mappings"): - mappings = getattr(module, "guardrail_translation_mappings") + mappings = attribute_or_none(module, "guardrail_translation_mappings") if isinstance(mappings, dict): discovered_mappings.update(mappings) verbose_logger.debug( diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 1818b406823..633ea73bb16 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -31,6 +31,10 @@ class A2AError(BaseLLMException): ) +def _message_role(message: Mapping[str, object]) -> object: + return message.get("role", "user") + + def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str: """ Convert OpenAI messages to a single prompt string for A2A agent. @@ -51,7 +55,7 @@ def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str: # Get role if isinstance(msg, BaseModel): - role = msg.model_dump().get("role", "user") + role = _message_role(msg.model_dump()) elif isinstance(msg, dict): role = msg.get("role", "user") else: diff --git a/litellm/llms/ai21/chat/transformation.py b/litellm/llms/ai21/chat/transformation.py index 2b8c36d0085..e78cc77c0fc 100644 --- a/litellm/llms/ai21/chat/transformation.py +++ b/litellm/llms/ai21/chat/transformation.py @@ -4,6 +4,8 @@ AI21 Chat Completions API this is OpenAI compatible - no translation needed / occurs """ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from ...openai_like.chat.transformation import OpenAILikeChatConfig @@ -41,7 +43,7 @@ class AI21ChatConfig(OpenAILikeChatConfig): tool_choice: str | None = None, user: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index 931038689cc..bb981983e98 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -121,7 +122,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): complete_url = complete_url.rstrip("/") # Strip /v1 suffix if present since IMAGE_GENERATION_ENDPOINT already includes v1 complete_url = complete_url.removesuffix("/v1") - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" + complete_url = join_api_base(complete_url, self.IMAGE_GENERATION_ENDPOINT) return complete_url def validate_environment( diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index dba7a7405ab..74f782b56e7 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final from aiohttp import ClientResponse from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, ModelResponse @@ -44,12 +45,7 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): Ensure - /v1/chat/completions is at the end of the url """ - if api_base is None: - api_base = "https://api.openai.com" - - if not api_base.endswith("/chat/completions"): - api_base += "/chat/completions" - return api_base + return ensure_api_route("https://api.openai.com" if api_base is None else api_base, "chat/completions") def validate_environment( self, diff --git a/litellm/llms/amazon_nova/chat/transformation.py b/litellm/llms/amazon_nova/chat/transformation.py index a93fbf1e933..4ef729bc76d 100644 --- a/litellm/llms/amazon_nova/chat/transformation.py +++ b/litellm/llms/amazon_nova/chat/transformation.py @@ -2,6 +2,8 @@ Translate from OpenAI's `/v1/chat/completions` to Amazon Nova's `/v1/chat/completions` """ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx @@ -38,7 +40,7 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): tools: list | None = None, reasoning_effort: list | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/anthropic/batches/handler.py b/litellm/llms/anthropic/batches/handler.py index 3f3716e0de6..a9d9f930b9e 100644 --- a/litellm/llms/anthropic/batches/handler.py +++ b/litellm/llms/anthropic/batches/handler.py @@ -132,7 +132,7 @@ class AnthropicBatchesHandler: max_retries: int | None, logging_obj: LiteLLMLoggingObj | None = None, litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment - ) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: + ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: """ Retrieve a batch from Anthropic. diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index b95b734c761..da6c2871289 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -7,7 +7,7 @@ import httpx from httpx import Headers, Response from typing_extensions import ReadOnly, TypedDict -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, ensure_api_route, join_api_base from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest @@ -106,10 +106,8 @@ class AnthropicBatchesConfig(BaseBatchesConfig): data: CreateBatchRequest, ) -> str: """Get the complete URL for batch creation request.""" - api_base = api_base or self.anthropic_model_info.get_api_base(api_base) - if not api_base.endswith("/v1/messages/batches"): - api_base = f"{api_base.rstrip('/')}/v1/messages/batches" - return api_base + api_base = self.anthropic_model_info.resolve_anthropic_api_base(api_base) + return ensure_api_route(api_base, "v1/messages/batches") def transform_create_batch_request( self, @@ -158,9 +156,9 @@ class AnthropicBatchesConfig(BaseBatchesConfig): Returns: Complete URL for Anthropic batch retrieval: {api_base}/v1/messages/batches/{batch_id} """ - api_base = api_base or self.anthropic_model_info.get_api_base(api_base) + api_base = self.anthropic_model_info.resolve_anthropic_api_base(api_base) encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") - return f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}" + return join_api_base(api_base, f"v1/messages/batches/{encoded_batch_id}") def transform_retrieve_batch_request( self, diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 4c6bded68ea..3ab21099d3a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1697,14 +1697,13 @@ class AnthropicMessagesHandler(BaseTranslation): if not line.startswith("data:"): return line try: - data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads( - line[len("data:") :].strip() - ) + data: Final = _DECODED_SSE_DATA.validate_python(json.loads(line[len("data:") :].strip())) except json.JSONDecodeError: return line - if not isinstance(data, dict): + event: Final = as_str_mapping(data) + if not isinstance(event, dict): return line - rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event) + rewritten: Final = _rewritten_event(event, rewrite_event) return line if rewritten is data else "data: " + json.dumps(rewritten) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index f0143f64be8..13ceea3e31b 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -856,11 +856,21 @@ class ModelResponseIterator: if reasoning_content is None: reasoning_content = "" elif "content" in content_block["delta"] and content_block["delta"].get("type") == "compaction_delta": - # Handle compaction delta + compaction_delta: Final = content_block["delta"]["content"] provider_specific_fields["compaction_delta"] = { "type": "compaction_delta", - "content": content_block["delta"]["content"], + "content": compaction_delta, } + if self.compaction_blocks and isinstance(compaction_delta, str): + streamed_summary: Final = self.compaction_blocks[-1].get("content") + self.compaction_blocks = [ + *self.compaction_blocks[:-1], + { + **self.compaction_blocks[-1], + "content": (streamed_summary if isinstance(streamed_summary, str) else "") + compaction_delta, + }, + ] + provider_specific_fields["compaction_blocks"] = self.compaction_blocks return ( text, diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 9a63afd5cc0..92e7c84031c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -831,6 +831,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "else": "conditional (else): {}", "not": "must not match: {}", } + schema_fields: Final[Mapping[str, object]] = MappingProxyType(schema) 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) @@ -845,7 +846,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): for field, label in constraint_labels.items(): if field not in schema or field not in unsupported_fields: continue - value = schema[field] + value = schema_fields[field] # A falsy boolean constraint (e.g. ``uniqueItems: false``) imposes no # real requirement, so don't add a misleading advisory note for it. if isinstance(value, bool) and not value: @@ -869,7 +870,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): drops_conflicting_type: Final = _enum_conflicts_with_declared_type(schema) - for key, value in schema.items(): + for key, value in schema_fields.items(): if key in unsupported_fields: continue if key == "type" and drops_conflicting_type: @@ -2385,7 +2386,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params["tools"] = tools ## Load Config - config: Final = litellm.AnthropicConfig.get_config(model=model) + config: Final[Mapping[str, object]] = MappingProxyType(litellm.AnthropicConfig.get_config(model=model)) for k, v in config.items(): if ( k not in optional_params diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 6f08d6b0d61..1f5e4e7de39 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -35,6 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import parts_of, ) from litellm.litellm_core_utils.reasoning_effort_utils import non_numeric_thinking_budget_error +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.wif import ( aget_anthropic_wif_token, anthropic_base_without_chat_suffix, @@ -468,7 +469,7 @@ def _fetch_anthropic_models_page( # which a read-only one cannot support query: Final = f"?after_id={quote(after_id)}" if after_id else "" response: Final = litellm.module_level_client.get( - url=f"{api_base}/v1/models{query}", + url=f"{join_api_base(api_base, 'v1/models')}{query}", headers=headers, follow_redirects=False, ) @@ -1406,7 +1407,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): return {**caller_headers, **anthropic_headers} @staticmethod - def get_api_base(api_base: str | None = None) -> str | None: + def resolve_anthropic_api_base(api_base: str | None = None) -> str: from litellm.secret_managers.main import get_secret_str return ( @@ -1416,6 +1417,10 @@ class AnthropicModelInfo(BaseLLMModelInfo): or "https://api.anthropic.com" ) + @staticmethod + def get_api_base(api_base: str | None = None) -> str | None: + return AnthropicModelInfo.resolve_anthropic_api_base(api_base) + @staticmethod def get_api_key(api_key: str | None = None) -> str | None: """An empty or whitespace-only key counts as unset: it can never authenticate anything, and @@ -1533,14 +1538,14 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_base: str | None, litellm_params: Mapping[str, object] | None, ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override - resolved_api_base: Final = AnthropicModelInfo.get_api_base(api_base) + resolved_api_base: Final = AnthropicModelInfo.resolve_anthropic_api_base(api_base) auth_header: Final = AnthropicModelInfo.get_auth_header( api_key, resolved_api_base, litellm_params=litellm_params, allow_workload_identity=config_allows_workload_identity(self), ) - if resolved_api_base is None or auth_header is None: + if auth_header is None: raise ValueError( "ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN (or workload " "identity federation via ANTHROPIC_FEDERATION_RULE_ID/ANTHROPIC_ORGANIZATION_ID/" diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index 73e94e85f04..5cc4e955d00 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -7,6 +7,7 @@ Litellm provider slug: `anthropic_text/` import json import time from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx @@ -112,7 +113,7 @@ class AnthropicTextConfig(BaseConfig): ) -> dict: prompt: Final = self._get_anthropic_text_prompt_from_messages(messages=messages, model=model) ## Load Config - config: Final = litellm.AnthropicTextConfig.get_config() + config: Final[Mapping[str, object]] = MappingProxyType(litellm.AnthropicTextConfig.get_config()) for k, v in config.items(): if ( k not in optional_params diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 592b5fb0f00..c2d87c641be 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import anthropic_system_blocks, split_leading_system_run, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers from litellm.llms.anthropic.wif import resolve_anthropic_base from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes @@ -317,7 +318,7 @@ class AnthropicCountTokensConfig: Returns: The endpoint URL for the CountTokens API """ - return resolve_anthropic_base(api_base) + "/v1/messages/count_tokens" + return join_api_base(resolve_anthropic_base(api_base), "v1/messages/count_tokens") def transform_request_to_count_tokens( self, diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index f1f9c204b58..279fb92cf6c 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -5,12 +5,13 @@ from collections.abc import Coroutine, Mapping from typing import Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( FileContentRequest, @@ -24,6 +25,8 @@ from litellm.types.utils import CallTypes, LlmProviders, ModelResponse from ..chat.transformation import AnthropicConfig from ..common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP, AnthropicModelInfo +_DECODED_JSON: Final = TypeAdapter(object) + class AnthropicFilesHandler: """ @@ -74,7 +77,7 @@ class AnthropicFilesHandler: batch_id = file_id # Get Anthropic API credentials - api_base = self.anthropic_model_info.get_api_base(api_base) + api_base = self.anthropic_model_info.resolve_anthropic_api_base(api_base) auth_header: Final = await self.anthropic_model_info.aget_auth_header( api_key, api_base, litellm_params=litellm_params, allow_workload_identity=True ) @@ -84,7 +87,7 @@ class AnthropicFilesHandler: # Construct the Anthropic batch results URL encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") - results_url: Final = f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}/results" + results_url: Final = join_api_base(api_base, f"v1/messages/batches/{encoded_batch_id}/results") # Prepare headers headers: Final = { @@ -194,7 +197,7 @@ class AnthropicFilesHandler: if not line.strip(): continue - anthropic_result: object = json.loads(line) + anthropic_result: object = _DECODED_JSON.validate_python(json.loads(line)) if not isinstance(anthropic_result, dict): raise TypeError( f"Anthropic batch result line is not a JSON object: {type(anthropic_result).__name__}" diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index 8433235dde4..db14d3366ff 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -22,7 +22,7 @@ from openai.types.file_deleted import FileDeleted from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, @@ -80,7 +80,7 @@ class AnthropicFilesConfig(BaseFilesConfig): stream: bool | None = None, ) -> str: api_base = AnthropicModelInfo.get_api_base(api_base) or ANTHROPIC_FILES_API_BASE - return f"{api_base.rstrip('/')}/v1/files" + return join_api_base(api_base, "v1/files") def get_error_class( self, @@ -226,7 +226,7 @@ class AnthropicFilesConfig(BaseFilesConfig): ) -> tuple[str, dict]: api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} + return join_api_base(api_base, f"v1/files/{encoded_file_id}"), {} def transform_retrieve_file_response( self, @@ -245,7 +245,7 @@ class AnthropicFilesConfig(BaseFilesConfig): ) -> tuple[str, dict]: api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} + return join_api_base(api_base, f"v1/files/{encoded_file_id}"), {} def transform_delete_file_response( self, @@ -264,7 +264,7 @@ class AnthropicFilesConfig(BaseFilesConfig): litellm_params: dict, ) -> tuple[str, dict]: api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE - url: Final = f"{api_base.rstrip('/')}/v1/files" + url: Final = join_api_base(api_base, "v1/files") params: Final[dict[str, str]] = {} if purpose: params["purpose"] = purpose @@ -299,8 +299,8 @@ class AnthropicFilesConfig(BaseFilesConfig): api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") if file_id.startswith(ANTHROPIC_MESSAGE_BATCH_ID_PREFIX): - return f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_file_id}/results", {} - return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}/content", {} + return join_api_base(api_base, f"v1/messages/batches/{encoded_file_id}/results"), {} + return join_api_base(api_base, f"v1/files/{encoded_file_id}/content"), {} def transform_file_content_response( self, diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 52b836c789a..4eb53f35a03 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -16,6 +16,7 @@ from typing import ( Protocol, cast, get_args, + runtime_checkable, ) from pydantic import TypeAdapter @@ -24,7 +25,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame @@ -65,6 +66,19 @@ class _ChunkStream(Protocol): def __aiter__(self) -> "AsyncIterator[ModelResponseStream]": ... +@runtime_checkable +class _AsyncCloseable(Protocol): + async def aclose(self) -> None: ... + + +@runtime_checkable +class _FailureLoggingWaiter(Protocol): + @property + def has_dispatched_failure_logging(self) -> bool: ... + + async def wait_for_failure_logging(self, timeout: float | None) -> None: ... + + def _optional_attr(obj: object, name: str) -> object: return getattr(obj, name, None) @@ -188,6 +202,22 @@ class _CombinedChunkSplitter: self._buffer: deque[ModelResponseStream] = deque() self._held_calls: Mapping[int, tuple[ModelResponseStream, _ToolCallPart]] = MappingProxyType({}) self._live_call_index: int | None = None + self._closed = False + + async def aclose(self) -> None: + if self._closed: + return + self._closed = True + if isinstance(self._stream, _AsyncCloseable): + await self._stream.aclose() + + @property + def has_dispatched_failure_logging(self) -> bool: + return isinstance(self._stream, _FailureLoggingWaiter) and self._stream.has_dispatched_failure_logging is True + + async def wait_for_failure_logging(self, timeout: float | None) -> None: + if isinstance(self._stream, _FailureLoggingWaiter): + await self._stream.wait_for_failure_logging(timeout=timeout) @property def chunks(self) -> "list[ModelResponseStream] | None": @@ -522,7 +552,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # Wrap the upstream stream so chunks that carry both content and a # finish_reason (fake-streamed providers) are split into two — see # _CombinedChunkSplitter. - super().__init__(_CombinedChunkSplitter(completion_stream)) + combined_chunk_splitter: Final = _CombinedChunkSplitter(completion_stream) + super().__init__(combined_chunk_splitter) + self._combined_chunk_splitter: Final = combined_chunk_splitter self.model = model self._message_id: str = f"msg_{uuid.uuid4()}" if litellm_logging_obj is not None: @@ -1220,6 +1252,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # For non-dict chunks, forward the original value unchanged yield chunk + async def aclose(self) -> None: + await self._combined_chunk_splitter.aclose() + async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]: """ Async version of anthropic_sse_wrapper. @@ -1244,7 +1279,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if provider_error is e: raise raise provider_error from e - if logging_obj is not None: + stream_has_dispatched_failure_logging: Final = self._combined_chunk_splitter.has_dispatched_failure_logging + if stream_has_dispatched_failure_logging: + await self._combined_chunk_splitter.wait_for_failure_logging(timeout=None) + elif logging_obj is not None: try: await logging_obj.dispatch_failure_handlers( exception=provider_error, @@ -1345,9 +1383,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # open a block: the emitted {"type": "thinking", "thinking": ""} gets # replayed as history and Anthropic rejects it (LIT-6357). A signed # entry opens the block so the client receives the replay signature. - thinking_blocks: Final = getattr(delta, "thinking_blocks", None) - if thinking_blocks and any( - isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) for b in thinking_blocks + if any( + isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) + for b in attribute_items(delta, "thinking_blocks") ): return False return True @@ -1459,4 +1497,7 @@ class AnthropicSSEStream(AsyncIterator[bytes]): return await self._byte_stream.__anext__() async def aclose(self) -> None: - await self._byte_stream.aclose() + try: + await self._byte_stream.aclose() + finally: + await self._anthropic_wrapper.aclose() diff --git a/litellm/llms/anthropic/pass_through/context_management/editors/clear_tool_uses.py b/litellm/llms/anthropic/pass_through/context_management/editors/clear_tool_uses.py index 935ce77ef87..d4430ca5ab9 100644 --- a/litellm/llms/anthropic/pass_through/context_management/editors/clear_tool_uses.py +++ b/litellm/llms/anthropic/pass_through/context_management/editors/clear_tool_uses.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from functools import reduce from itertools import chain from types import MappingProxyType -from typing import Annotated, Any, Final, Literal, TypeVar, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeVar, cast from pydantic import BaseModel, ConfigDict, Field, StrictInt, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -32,6 +32,9 @@ from ..constants import ( from ..errors import AnthropicContextManagementError from ..placeholders import build_cleared_tool_result_content +if TYPE_CHECKING: + from litellm.types.llms.openai import ChatCompletionToolParam + T = TypeVar("T") # api.anthropic.com, measured 2026-10-19: text results of 1 to 46 bytes are never cleared or counted toward keep _TINY_RESULT_BYTES: Final = 46 @@ -216,7 +219,7 @@ def _trigger_met( current_tokens: Final = litellm.token_counter( model=model, messages=messages, - tools=cast(Any, tools), + tools=cast("list[ChatCompletionToolParam] | None", tools), ) verbose_logger.debug("context_management polyfill: current_tokens: %s", current_tokens) verbose_logger.debug("context_management polyfill: threshold: %s", threshold) @@ -366,9 +369,13 @@ def apply_clear_tool_uses_20250919( tokens_before: Final = ( counted_tokens if counted_tokens is not None - else litellm.token_counter(model=model, messages=messages, tools=cast(Any, tools)) + else litellm.token_counter( + model=model, messages=messages, tools=cast("list[ChatCompletionToolParam] | None", tools) + ) + ) + tokens_after: Final = litellm.token_counter( + model=model, messages=final_messages, tools=cast("list[ChatCompletionToolParam] | None", tools) ) - tokens_after: Final = litellm.token_counter(model=model, messages=final_messages, tools=cast(Any, tools)) cleared_input_tokens: Final = max(tokens_before - tokens_after, 0) return final_messages, { diff --git a/litellm/llms/anthropic/pass_through/context_management/placeholders.py b/litellm/llms/anthropic/pass_through/context_management/placeholders.py index b3d3529f105..8a0eb6d48e1 100644 --- a/litellm/llms/anthropic/pass_through/context_management/placeholders.py +++ b/litellm/llms/anthropic/pass_through/context_management/placeholders.py @@ -1,12 +1,10 @@ """Placeholder content for cleared ``tool_result`` blocks (string or block list).""" -from typing import Any - from .constants import CLEARED_TOOL_RESULT_PLACEHOLDER def build_cleared_tool_result_content( - original_content: Any, + original_content: object, ) -> str | list[dict]: """Return a string or single text block list, matching ``original_content`` shape.""" if isinstance(original_content, list): diff --git a/litellm/llms/anthropic/pass_through/context_management/result.py b/litellm/llms/anthropic/pass_through/context_management/result.py index 5b88fe30871..4713b8d79c7 100644 --- a/litellm/llms/anthropic/pass_through/context_management/result.py +++ b/litellm/llms/anthropic/pass_through/context_management/result.py @@ -6,7 +6,7 @@ attach ``iterations`` to ``usage``. """ from dataclasses import dataclass, field -from typing import Any, Final +from typing import Final from litellm.types.llms.anthropic import ( AppliedEdit, @@ -19,8 +19,8 @@ from .constants import COMPACT_EDIT_TYPE @dataclass class PolyfillResult: - messages: list[dict[str, Any]] - system: str | list[dict[str, Any]] | None + messages: list[dict[str, object]] + system: str | list[dict[str, object]] | None applied_edits: list[AppliedEdit] = field(default_factory=list) compaction_block: CompactionBlock | None = None iterations_usage: list[UsageIteration] | None = None diff --git a/litellm/llms/anthropic/pass_through/convertible_blocks.py b/litellm/llms/anthropic/pass_through/convertible_blocks.py index fb17910fb97..9d651f8d721 100644 --- a/litellm/llms/anthropic/pass_through/convertible_blocks.py +++ b/litellm/llms/anthropic/pass_through/convertible_blocks.py @@ -3,7 +3,6 @@ from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial from itertools import chain, groupby -from operator import itemgetter from types import MappingProxyType from typing import Final, TypeAlias @@ -279,6 +278,14 @@ def _with_kept_parts(block: object, tool_result: Mapping[str, object], parts: _K return {**tool_result, "content": [_restored_part(part, parts.get(index)) for index, part in enumerate(blocks)]} +def _entry_tool_use_id(entry: tuple[str, _KeptParts]) -> str: + return entry[0] + + +def _entry_kept_parts(entry: tuple[str, _KeptParts]) -> _KeptParts: + return entry[1] + + def _restored_block(block: object, kept: Mapping[str, Iterator[_KeptParts]]) -> object: match as_str_mapping(block): case {"type": "tool_result", "tool_use_id": str(tool_use_id)} as tool_result if ( @@ -293,13 +300,15 @@ def _restored_block(block: object, kept: Mapping[str, Iterator[_KeptParts]]) -> def keep_non_string_tool_result_texts(rewritten: Sequence[object], original: object) -> Sequence[object]: blocks: Final = chain.from_iterable(map(_user_content_blocks, _list_items(original))) - by_id: Final = tuple(sorted(chain.from_iterable(map(_tool_result_non_string_texts, blocks)), key=itemgetter(0))) + by_id: Final = tuple( + sorted(chain.from_iterable(map(_tool_result_non_string_texts, blocks)), key=_entry_tool_use_id) + ) if not any(parts for _, parts in by_id): return rewritten kept_in_order: Final = MappingProxyType( { - tool_use_id: iter(tuple(map(itemgetter(1), group))) - for tool_use_id, group in groupby(by_id, key=itemgetter(0)) + tool_use_id: iter(tuple(map(_entry_kept_parts, group))) + for tool_use_id, group in groupby(by_id, key=_entry_tool_use_id) } ) return [_with_user_blocks(message, partial(_restored_block, kept=kept_in_order)) for message in rewritten] 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 8d1b358c8f1..2b6161d98f6 100644 --- a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py @@ -9,8 +9,8 @@ the LLM doesn't make a tool call, and we need to return a stream to the user. """ import json -from collections.abc import Mapping -from typing import Any, Final, cast +from collections.abc import Mapping, Sequence +from typing import Final, cast from litellm.llms.anthropic.pass_through.messages.utils import anthropic_message_delta from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -136,8 +136,7 @@ class FakeAnthropicMessagesStreamIterator: """Convert the non-streaming response to streaming chunks""" chunks: Final = [] - # Cast response to dict for easier access - response_dict: Final = cast(dict[str, Any], self.response) + response_dict: Final[Mapping[str, object]] = self.response # 1. message_start event usage: Final = self.response.get("usage") @@ -160,9 +159,8 @@ class FakeAnthropicMessagesStreamIterator: chunks.append(f"event: message_start\ndata: {json.dumps(message_start)}\n\n".encode()) # 2-4. For each content block, send start/delta/stop events - content_blocks: Final = response_dict.get("content", []) - for index, block in enumerate(content_blocks): - block_dict = cast(dict[str, Any], block) + content_blocks: Final = cast(Sequence[Mapping[str, object]], response_dict.get("content", [])) + for index, block_dict in enumerate(content_blocks): chunks.extend(self._create_content_block_chunks(block_dict, index)) # 5. message_delta event (with final usage and stop_reason) diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index a71cd72b768..35db54147c5 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -9,6 +9,7 @@ from litellm.anthropic_beta_headers_manager import get_unsupported_headers, is_b from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.pass_through.convertible_blocks import drop_non_string_tool_result_texts from litellm.llms.anthropic.pass_through.safeguards import proxy_serves_safeguards from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -74,6 +75,12 @@ def messages_carry_output_config(messages: Sequence[object]) -> bool: _messages_carry_output_config = messages_carry_output_config +def _ensure_stop_sequence(response: AnthropicMessagesResponse) -> AnthropicMessagesResponse: + if "stop_sequence" not in response: + response["stop_sequence"] = None # rebind-ok: normalize the request-owned response in place + return response + + class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): _workload_identity_eligible: ClassVar[bool] = True @@ -287,10 +294,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com" - if not api_base.endswith("/v1/messages"): - api_base = f"{api_base}/v1/messages" - return api_base + return ensure_api_route(AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com", "v1/messages") def validate_anthropic_messages_environment( self, @@ -717,14 +721,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> AnthropicMessagesResponse: - """ - No transformation is needed for Anthropic messages, since we want the response in the Anthropic /v1/messages API spec - """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[AnthropicMessagesResponse] = raw_response.json() except Exception: raise AnthropicError(message=raw_response.text, status_code=raw_response.status_code) - return AnthropicMessagesResponse(**raw_response_json) + response: Final = AnthropicMessagesResponse(**raw_response_json) + return _ensure_stop_sequence(response) def get_async_streaming_response_iterator( self, diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py index bd0f57c3ed9..48fc987a95a 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py @@ -242,7 +242,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def _summary_part_text(part: object) -> str: if isinstance(part, Mapping): - mapping: Final = cast(Mapping[str, Any], part) # cast-ok: summary parts are untyped provider json + mapping: Final = cast(Mapping[str, object], part) # cast-ok: summary parts are untyped provider json return str(mapping.get("text") or "") return str(getattr(part, "text", None) or "") @@ -673,7 +673,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: output_config = anthropic_request.get("output_config") reasoning: Final = self.translate_thinking_to_reasoning( thinking, - output_config=cast(dict[str, Any] | None, output_config), + output_config=cast(dict[str, object] | None, output_config), auto_summary_level=auto_summary_level, drop_params=drop_params, ) diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 6d419a56889..8ea9d192c24 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -9,7 +9,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.skills.transformation import ( BaseSkillsAPIConfig, LiteLLMLoggingObj, @@ -75,13 +75,11 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): """Get complete URL for Anthropic Skills API""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo - if api_base is None: - api_base = AnthropicModelInfo.get_api_base() - + resolved_api_base: Final = AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com" if skill_id: encoded_skill_id: Final = encode_url_path_segment(skill_id, field_name="skill_id") - return f"{api_base}/v1/skills/{encoded_skill_id}" - return f"{api_base}/v1/{endpoint}" + return join_api_base(resolved_api_base, f"v1/skills/{encoded_skill_id}") + return join_api_base(resolved_api_base, f"v1/{endpoint}") def transform_create_skill_request( self, diff --git a/litellm/llms/anthropic/wif.py b/litellm/llms/anthropic/wif.py index b47c4f28eb2..ccb81c1244f 100644 --- a/litellm/llms/anthropic/wif.py +++ b/litellm/llms/anthropic/wif.py @@ -15,6 +15,7 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.auth.client_credentials import keycloak_assertion_source from litellm.llms.base_llm.auth.identity_source import ( AnthropicIdentitySourceKind, @@ -324,7 +325,7 @@ def anthropic_internal_issuer_jwks(credential_values: Mapping[str, object]) -> A def build_anthropic_wif_spec(params: AnthropicWifParams, api_base: str) -> TokenExchangeSpec: return TokenExchangeSpec( - token_url=api_base.rstrip("/") + ANTHROPIC_TOKEN_EXCHANGE_PATH, + token_url=join_api_base(api_base, ANTHROPIC_TOKEN_EXCHANGE_PATH), assertion_ref=params.assertion_ref, assertion_field="assertion", static_body=MappingProxyType( diff --git a/litellm/llms/apiserpent/search/transformation.py b/litellm/llms/apiserpent/search/transformation.py index bb8abff7272..faf1903ffe8 100644 --- a/litellm/llms/apiserpent/search/transformation.py +++ b/litellm/llms/apiserpent/search/transformation.py @@ -14,6 +14,7 @@ from urllib.parse import urlencode import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.apiserpent.search.defaults import ( DEEP_SEARCH_PATH, NUM_MAX, @@ -77,13 +78,11 @@ class APISerpentSearchConfig(BaseSearchConfig): Build the search URL. APISerpent uses GET, so the transformed request is serialized into the query string. The endpoint path (quick vs deep) is always applied; an ``api_base`` / ``APISERPENT_API_BASE`` override only - changes the host. The ``endswith`` guard keeps this idempotent, since the + changes the host. ``ensure_api_route`` keeps this idempotent, since the handler re-invokes this method with the already-resolved URL as api_base. """ - base = (api_base or get_secret_str("APISERPENT_API_BASE") or APISERPENT_BASE).rstrip("/") path: Final = DEEP_SEARCH_PATH if self._is_deep_search(optional_params) else QUICK_SEARCH_PATH - if not base.endswith(path): - base = f"{base}{path}" + base: Final = ensure_api_route(api_base or get_secret_str("APISERPENT_API_BASE") or APISERPENT_BASE, path) if data and isinstance(data, dict) and APISERPENT_PARAMS_KEY in data: query_string: Final = urlencode(data[APISERPENT_PARAMS_KEY], doseq=True) diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 01fec02c310..325c9fd364d 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -178,7 +178,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): else: mapped_params["output_format"] = self.DEFAULT_OUTPUT_FORMAT - # Extract engine from model name (e.g., "aws_polly/neural" -> "neural") + # model arrives without the "aws_polly/" prefix, e.g. "standard" engine: Final = self._extract_engine_from_model(model) mapped_params["engine"] = engine @@ -196,18 +196,17 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): """ Extract engine from model name. + speech() strips the "aws_polly/" prefix before the config is called, + so the bare engine name is what normally arrives here. + Examples: - - aws_polly/neural -> neural - - aws_polly/standard -> standard + - standard -> standard - aws_polly/long-form -> long-form - aws_polly -> neural (default) """ - if "/" in model: - parts: Final = model.split("/") - if len(parts) >= 2: - engine: Final = parts[1].lower() - if engine in self.VALID_ENGINES: - return engine + engine: Final = model.rsplit("/", 1)[-1].lower() + if engine in self.VALID_ENGINES: + return engine return self.DEFAULT_ENGINE def validate_environment( diff --git a/litellm/llms/azure/audio_transcription/transformation.py b/litellm/llms/azure/audio_transcription/transformation.py index 858f122f65a..a5fde7f7158 100644 --- a/litellm/llms/azure/audio_transcription/transformation.py +++ b/litellm/llms/azure/audio_transcription/transformation.py @@ -13,6 +13,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -110,7 +111,7 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): "language": optional_params.get("language", self.DEFAULT_LANGUAGE), "format": self._get_azure_response_format(optional_params.get("response_format")), } - return f"{base_url}{self.STT_ENDPOINT_PATH}?{urlencode(query_params)}" + return f"{join_api_base(base_url, self.STT_ENDPOINT_PATH)}?{urlencode(query_params)}" def transform_audio_transcription_request( self, @@ -153,7 +154,6 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) def _resolve_stt_base_url(self, api_base: str) -> str: - api_base = api_base.rstrip("/") parsed_url: Final = urlparse(api_base) hostname: Final = parsed_url.hostname or "" diff --git a/litellm/llms/azure/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index 99ebd2875a7..df8a0614665 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -1,5 +1,5 @@ from collections.abc import Coroutine -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from openai import AsyncAzureOpenAI, AzureOpenAI from pydantic import BaseModel @@ -38,7 +38,7 @@ class AzureAudioTranscription(AzureChatCompletion): atranscription: bool = False, litellm_params: dict | None = None, custom_llm_provider: str = "azure", - ) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]: + ) -> TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]: data: Final = {"model": model, "file": audio_file, **optional_params} if atranscription is True: diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 67b25bf7d12..02f1ba3472a 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -13,7 +13,7 @@ from openai import ( BadRequestError, OpenAI, ) -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES @@ -59,6 +59,8 @@ from .image_generation import ( ) from .image_generation.http_utils import azure_deployment_image_generation_json_body +_DECODED_JSON: Final = TypeAdapter(object) + class AzureOpenAIAssistantsAPIConfig: """ @@ -961,7 +963,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): body=error_data, ) - result: Final = response.json()["result"] + result: Final = _DECODED_JSON.validate_python(response.json()["result"]) return httpx.Response( status_code=200, headers=response.headers, @@ -1078,7 +1080,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): body=error_data, ) - result: Final = response.json()["result"] + result: Final = _DECODED_JSON.validate_python(response.json()["result"]) return httpx.Response( status_code=200, headers=response.headers, @@ -1358,7 +1360,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except AzureOpenAIError as e: raise e except Exception as e: - error_code: Final = getattr(e, "status_code", None) + error_code: Final[object] = getattr(e, "status_code", None) if error_code is not None: raise AzureOpenAIError(status_code=error_code, message=str(e)) else: diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 52e3cdf380e..23c5aa35f4b 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -18,6 +18,7 @@ import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.constants import DEFAULT_MAX_RETRIES +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient @@ -616,7 +617,7 @@ class BaseAzureLLM(BaseOpenAILLM): v1_params: Final[_AzureV1ClientParams] = { "api_key": v1_api_key, - "base_url": f"{api_base}/openai/v1/", + "base_url": f"{ensure_api_route(api_base or '', 'openai/v1')}/", } if "timeout" in azure_client_params: v1_params["timeout"] = azure_client_params["timeout"] diff --git a/litellm/llms/azure/exception_mapping.py b/litellm/llms/azure/exception_mapping.py index dcc3601cd8e..f7146a36384 100644 --- a/litellm/llms/azure/exception_mapping.py +++ b/litellm/llms/azure/exception_mapping.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from litellm.exceptions import ContentPolicyViolationError @@ -54,7 +54,7 @@ class AzureOpenAIExceptionMapping: @staticmethod def _extract_azure_error( original_exception: Exception, - ) -> tuple[dict[str, object], dict | None]: + ) -> tuple[dict[str, object], object]: """Extract Azure OpenAI error payload and inner error details. Azure error formats can vary by endpoint/version. Common shapes: @@ -67,13 +67,13 @@ class AzureOpenAIExceptionMapping: return {}, None # Some SDKs place the payload under "error". - azure_error: dict[str, Any] + azure_error: dict[str, object] if isinstance(body_dict.get("error"), dict): azure_error = body_dict.get("error", {}) else: azure_error = body_dict - inner_error: Final = ( + inner_error: Final[object] = ( azure_error.get("inner_error") or azure_error.get("innererror") or body_dict.get("innererror") diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 75f96bfdf63..3d4576fd7ac 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -10,6 +10,7 @@ from typing import Final, Protocol, cast from litellm._logging import verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -136,7 +137,8 @@ class AzureOpenAIRealtime(AzureChatCompletion): query_parts.append(urlencode({"intent": intent})) qs: Final = "&".join(query_parts) - return f"{api_base}{path}?{qs}" if qs else f"{api_base}{path}" + endpoint_url: Final = join_api_base(api_base, path) + return f"{endpoint_url}?{qs}" if qs else endpoint_url def construct_url( self, diff --git a/litellm/llms/azure/realtime/http_transformation.py b/litellm/llms/azure/realtime/http_transformation.py index 542a6a8c02f..545c3348182 100644 --- a/litellm/llms/azure/realtime/http_transformation.py +++ b/litellm/llms/azure/realtime/http_transformation.py @@ -3,6 +3,7 @@ from typing import Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig from litellm.secret_managers.main import get_secret_str @@ -39,9 +40,9 @@ class AzureRealtimeHTTPConfig(BaseRealtimeHTTPConfig): return _azure_realtime_v1_url(self.get_api_base(api_base), "calls") def get_transcription_session_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: - base: Final = self.get_api_base(api_base).rstrip("/") + base: Final = self.get_api_base(api_base) version: Final = api_version or get_secret_str("AZURE_API_VERSION") or "2024-12-17" - return f"{base}/openai/realtime/transcription_sessions?api-version={version}" + return f"{join_api_base(base, 'openai/realtime/transcription_sessions')}?api-version={version}" def get_realtime_calls_headers(self, ephemeral_key: str) -> dict[str, str]: return { diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 74e65572809..6d904503df8 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -3,6 +3,7 @@ from collections.abc import Mapping from copy import deepcopy from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal +from urllib.parse import urlsplit import httpx from openai.types.responses import ResponseReasoningItem @@ -94,7 +95,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def get_stripped_model_name(self, model: str) -> str: return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "") - def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]: + def _handle_reasoning_item(self, item: dict[str, Any]) -> Mapping[str, object]: """ Handle reasoning items to filter out the status field. Issue: https://github.com/BerriAI/litellm/issues/13484 @@ -211,9 +212,21 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2024-05-01-preview" """ resolved_api_base: Final = BaseAzureLLM.resolve_azure_api_base(api_base) + responses_params: Final = _with_responses_api_version(resolved_api_base, litellm_params) + parsed_base: Final = urlsplit(resolved_api_base or "") + path: Final = parsed_base.path.rstrip("/") + if path.endswith(("/openai/v1", "/openai/v1/responses")): + responses_path: Final = f"{path}/responses" if path.endswith("/openai/v1") else path + return BaseAzureLLM.get_base_azure_url( + api_base=parsed_base._replace(path=responses_path).geturl(), + litellm_params=responses_params, + route="/openai/v1/responses", + default_api_version=AZURE_DEFAULT_RESPONSES_API_VERSION, + ) + return BaseAzureLLM.get_base_azure_url( api_base=resolved_api_base, - litellm_params=_with_responses_api_version(resolved_api_base, litellm_params), + litellm_params=responses_params, route="/openai/responses", default_api_version=AZURE_DEFAULT_RESPONSES_API_VERSION, ) diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py index f2bb2ecc68c..f7cdf88ade0 100644 --- a/litellm/llms/azure/search/transformation.py +++ b/litellm/llms/azure/search/transformation.py @@ -35,6 +35,7 @@ from typing import TYPE_CHECKING, Final, Literal import httpx from pydantic import ConfigDict, ValidationError +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( @@ -340,10 +341,7 @@ class BingGroundingSearchConfig(BaseSearchConfig): f"{PROJECT_ENDPOINT_ENV} is not set. Set it to your Microsoft Foundry project " f"endpoint, e.g. https://.services.ai.azure.com/api/projects/." ) - trimmed: Final = resolved_base.rstrip("/") - if trimmed.endswith(_RESPONSES_PATH): - return trimmed - return f"{trimmed}{_RESPONSES_PATH}" + return ensure_api_route(resolved_base, _RESPONSES_PATH) def transform_search_request( self, diff --git a/litellm/llms/azure/text_to_speech/transformation.py b/litellm/llms/azure/text_to_speech/transformation.py index eed7a3178ca..acae650ed75 100644 --- a/litellm/llms/azure/text_to_speech/transformation.py +++ b/litellm/llms/azure/text_to_speech/transformation.py @@ -11,6 +11,7 @@ from urllib.parse import urlparse import httpx import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -329,8 +330,6 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): f"https://{{region}}.{self.TTS_SPEECH_DOMAIN}" ) - # Remove trailing slash and parse URL - api_base = api_base.rstrip("/") parsed_url: Final = urlparse(api_base) hostname: Final = parsed_url.hostname or "" @@ -339,14 +338,10 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): region: Final = self._extract_region_from_hostname(hostname=hostname, domain=self.COGNITIVE_SERVICES_DOMAIN) return self._build_tts_url(region=region) - # Check if it's already a TTS endpoint if self._is_tts_endpoint(hostname=hostname): - if not api_base.endswith(self.TTS_ENDPOINT_PATH): - return f"{api_base}{self.TTS_ENDPOINT_PATH}" - return api_base + return ensure_api_route(api_base, self.TTS_ENDPOINT_PATH) - # Assume it's a custom endpoint, append the path - return f"{api_base}{self.TTS_ENDPOINT_PATH}" + return join_api_base(api_base, self.TTS_ENDPOINT_PATH) def _is_cognitive_services_endpoint(self, hostname: str) -> bool: """Check if hostname is a Cognitive Services endpoint""" diff --git a/litellm/llms/azure/vector_stores/transformation.py b/litellm/llms/azure/vector_stores/transformation.py index 9510c021c7d..5a2961cd550 100644 --- a/litellm/llms/azure/vector_stores/transformation.py +++ b/litellm/llms/azure/vector_stores/transformation.py @@ -1,3 +1,5 @@ +from typing import Final + from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig from litellm.types.router import GenericLiteLLMParams @@ -15,5 +17,10 @@ class AzureOpenAIVectorStoreConfig(OpenAIVectorStoreConfig): route="/openai/vector_stores", ) + def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> str: + path, _, query = api_base.partition("?") + url: Final = f"{path}/{encoded_vector_store_id}" + return f"{url}?{query}" if query else url + def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: return BaseAzureLLM.base_validate_azure_environment(headers=headers, litellm_params=litellm_params) diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 0b0548dd3a2..94f1399f01b 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -30,7 +30,7 @@ from typing_extensions import ReadOnly from litellm._logging import verbose_logger from litellm.litellm_core_utils.hidden_params import set_hidden_param -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.azure_ai.agents.transformation import ( AzureAIAgentsConfig, AzureAIAgentsError, @@ -130,28 +130,29 @@ class AzureAIAgentsHandler: # See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart # ------------------------------------------------------------------------- def _build_thread_url(self, api_base: str, api_version: str) -> str: - return f"{api_base}/threads?api-version={api_version}" + return f"{join_api_base(api_base, 'threads')}?api-version={api_version}" def _build_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") - return f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" + return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/messages')}?api-version={api_version}" def _build_runs_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") - return f"{api_base}/threads/{encoded_thread_id}/runs?api-version={api_version}" + return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/runs')}?api-version={api_version}" def _build_run_status_url(self, api_base: str, thread_id: str, run_id: str, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") - return f"{api_base}/threads/{encoded_thread_id}/runs/{encoded_run_id}?api-version={api_version}" + run_status_route: Final = f"threads/{encoded_thread_id}/runs/{encoded_run_id}" + return f"{join_api_base(api_base, run_status_route)}?api-version={api_version}" def _build_list_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") - return f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" + return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/messages')}?api-version={api_version}" def _build_create_thread_and_run_url(self, api_base: str, api_version: str) -> str: """URL for the create-thread-and-run endpoint (supports streaming).""" - return f"{api_base}/threads/runs?api-version={api_version}" + return f"{join_api_base(api_base, 'threads/runs')}?api-version={api_version}" # ------------------------------------------------------------------------- # Response Helpers @@ -281,8 +282,6 @@ class AzureAIAgentsHandler: api_version: Final = optional_params.get("api_version", self.config.DEFAULT_API_VERSION) agent_id: Final = self.config.get_agent_id(model, optional_params) thread_id: Final = optional_params.get("thread_id") - api_base = api_base.rstrip("/") - verbose_logger.debug("Azure AI Agents completion - api_base: %s, agent_id: %s", api_base, agent_id) return headers, api_version, agent_id, thread_id, api_base diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index 5946a1f4480..6e3cfce5eb4 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -5,8 +5,10 @@ Extends the base Anthropic CountTokens transformation with Azure authentication. """ from typing import Final +from urllib.parse import urlsplit from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -76,14 +78,8 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): Returns: The endpoint URL for the CountTokens API """ - # Azure AI Anthropic endpoint format: - # https://.services.ai.azure.com/anthropic/v1/messages/count_tokens - api_base = api_base.rstrip("/") - - # Ensure the URL has /anthropic path - if not api_base.endswith("/anthropic"): - if "/anthropic" not in api_base: - api_base = f"{api_base}/anthropic" - - # Add the count_tokens path - return f"{api_base}/v1/messages/count_tokens" + base: Final = api_base.rstrip("/") + if base.endswith("/v1/messages"): + return join_api_base(base, "count_tokens") + anthropic_base: Final = base if "/anthropic" in urlsplit(base).path else join_api_base(base, "anthropic") + return join_api_base(anthropic_base, "v1/messages/count_tokens") diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 91c75661688..2aacc5771d0 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -9,6 +9,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.azure_ai.anthropic.transformation import azure_anthropic_messages_url from litellm.types.router import GenericLiteLLMParams @@ -96,25 +97,7 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): "Expected format: https://.services.ai.azure.com/anthropic" ) - # Ensure the URL ends with /v1/messages - api_base = api_base.rstrip("/") - if api_base.endswith("/v1/messages") or api_base.endswith("/anthropic/v1/messages"): - # Already correct - pass - else: - # Check if /anthropic is already in the path - if "/anthropic" in api_base: - # /anthropic exists, ensure we end with /anthropic/v1/messages - # Extract the base URL up to and including /anthropic - parts: Final = api_base.split("/anthropic", 1) - api_base = parts[0] + "/anthropic" - else: - # /anthropic not in path, add it - api_base = api_base + "/anthropic" - # Add /v1/messages - api_base = api_base + "/v1/messages" - - return api_base + return azure_anthropic_messages_url(api_base) def _remove_scope_from_cache_control(self, anthropic_messages_request: dict) -> None: """ diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index 18a4eeb7c0a..c75213c3e2c 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -4,12 +4,23 @@ Azure Anthropic transformation config - extends AnthropicConfig with Azure authe from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams +def azure_anthropic_messages_url(api_base: str) -> str: + base: Final = api_base.rstrip("/") + if base.endswith("/v1/messages"): + return base + path_start: Final = base.find("/", base.find("://") + len("://")) if "://" in base else 0 + anthropic_start: Final = base.find("/anthropic", path_start) if path_start >= 0 else -1 + anthropic_root: Final = base[:anthropic_start] if anthropic_start >= 0 else base + return join_api_base(anthropic_root, "anthropic/v1/messages") + + def _promote_extra_body_to_optional_params(optional_params: dict) -> None: """Promote anthropic-native passthrough keys out of ``extra_body``. diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 0fa54cccbc9..3b541310d28 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -12,6 +12,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( audio_or_image_in_message_content, convert_content_list_to_str, + drop_tool_reference_parts_from_tool_messages, filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM @@ -250,7 +251,7 @@ class AzureAIStudioConfig(OpenAIConfig): Operates on a deep copy so the caller's messages keep their thinking blocks and provider metadata, which a fallback to another provider still needs. """ - stripped_messages: Final = copy.deepcopy(messages) + stripped_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages)) for message in stripped_messages: message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: @@ -396,7 +397,7 @@ class AzureAIStudioConfig(OpenAIConfig): return data def _drop_tool_level_extra_fields(self, request_data: dict, error_text: str) -> dict: - fields_to_drop: Final = set(re.findall(r"tools\[\d+\]\.([\w-]+)", error_text)) + fields_to_drop: Final[Collection[str]] = set(re.findall(r"tools\[\d+\]\.([\w-]+)", error_text)) tools: Final = request_data.get("tools") if fields_to_drop and isinstance(tools, list): for tool in tools: diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 57e6ca51ade..dc3bc6bdc5e 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -5,6 +5,7 @@ from urllib.parse import urlsplit, urlunsplit from openai import OpenAI import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -74,13 +75,18 @@ class AzureAIEmbedding(OpenAIChatCompletion): api_base: str | None, client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> EmbeddingResponse: + if api_base is None: + raise ValueError( + "api_base is None. Please set AZURE_AI_API_BASE or dynamically via `api_base` param, " + "to make the request." + ) if client is None or not isinstance(client, AsyncHTTPHandler): client = get_async_httpx_client( llm_provider=litellm.LlmProviders.AZURE_AI, params={"timeout": timeout}, ) - url: Final = f"{api_base}/images/embeddings" + url: Final = join_api_base(api_base, "images/embeddings") response: Final = await client.post( url=url, @@ -123,7 +129,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): if client is None or not isinstance(client, HTTPHandler): client = HTTPHandler(timeout=timeout, concurrent_limit=1) - url: Final = f"{api_base}/images/embeddings" + url: Final = join_api_base(api_base, "images/embeddings") response: Final = client.post( url=url, diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 8b4fa78cdf4..6a7bebd9f6b 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -3,6 +3,7 @@ from types import MappingProxyType from typing import Final from litellm.exceptions import BadRequestError, UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.image_generation import GPTImageGenerationConfig from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams @@ -52,7 +53,8 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): return f"{api_base}?api-version={api_version}" provider_model_path: Final = AzureFoundryFluxImageGenerationConfig.get_flux2_provider_model_path(model) - return f"{api_base}/providers/blackforestlabs/v1/{provider_model_path}?api-version={api_version}" + flux_route: Final = f"providers/blackforestlabs/v1/{provider_model_path}" + return f"{join_api_base(api_base, flux_route)}?api-version={api_version}" @staticmethod def is_flux2_model(model: str) -> bool: diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 18b4b6f456a..8bb5d3e6b8d 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -41,9 +42,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if "/mai/" in path: prefix, _, _ = path.partition("/images/") - path = f"{prefix}/images/generations" + path = join_api_base(prefix, "images/generations") else: - path = f"{path}/mai/v1/images/generations" + path = join_api_base(path, "mai/v1/images/generations") if separator: return f"{path}?{query}" @@ -63,9 +64,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if "/mai/" in path: prefix, _, _ = path.partition("/images/") - path = f"{prefix}/images/edits" + path = join_api_base(prefix, "images/edits") else: - path = f"{path}/mai/v1/images/edits" + path = join_api_base(path, "mai/v1/images/edits") if separator: return f"{path}?{query}" diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 7638229bc32..a4777be146f 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, @@ -172,7 +173,8 @@ class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM litellm_logging_obj.model_call_details["input"] = query_text litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") litellm_logging_obj.model_call_details["top_k"] = top_k - return f"{api_base}/indexes/{vector_store_id}/docs/search?api-version=2024-07-01", { + search_route: Final = f"indexes/{vector_store_id}/docs/search" + return f"{join_api_base(api_base, search_route)}?api-version=2024-07-01", { "search": "*", "vectorQueries": [{"vector": query_vector, "fields": vector_field, "kind": "vector", "k": top_k}], "select": "id,content", diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index b2eb4cde5ce..207a4705c92 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -94,9 +94,10 @@ class BaseConfig(ABC): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("_") and not isinstance( v, @@ -437,6 +438,7 @@ class BaseConfig(ABC): sync_stream: bool, json_mode: bool | None, optional_params: Mapping[str, object], + logging_obj: LiteLLMLoggingObj | None = None, ) -> object: if sync_stream: return self.get_model_response_iterator( @@ -513,7 +515,7 @@ class BaseConfig(ABC): def apply_assembled_streaming_response_metadata( self, response: "ModelResponse", - chunks: list[Any], + chunks: list[object], ) -> None: """Hook for providers to merge chunk metadata into assembled streaming responses.""" return diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index d74257a4a9d..14c6a2d2a66 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -12,6 +12,7 @@ from pydantic_core import ErrorDetails from typing_extensions import assert_never import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.decisions import ( @@ -375,12 +376,14 @@ def _ir_answer(question: DecisionsIRQuestion, answer: DecisionAnswer | None) -> def systemone_response_to_ir(response: DecisionsResponse, request: DecisionsIRRequest) -> DecisionsIRResponse: - usage: Final = response.usage or DecisionsUsage() + usage: Final = response.usage keyed_questions: Final = zip(systemone_keys(request.questions), request.questions, strict=True) return DecisionsIRResponse( model=response.model, answers=tuple(_ir_answer(question, response.answers.get(key)) for key, question in keyed_questions), - usage=DecisionsIRUsage( + usage=None + if usage is None + else DecisionsIRUsage( input_tokens=usage.input_tokens, output_tokens=usage.output_tokens, cached_tokens=usage.cached_tokens, @@ -484,7 +487,9 @@ def ir_to_systemone_response(response: DecisionsIRResponse, request: DecisionsIR **response.extra, "model": response.model, "answers": {key: answer for key, answer in keyed_answers if answer is not None}, - "usage": DecisionsUsage.model_validate( + "usage": None + if response.usage is None + else DecisionsUsage.model_validate( { **response.usage.extra, "input_tokens": response.usage.input_tokens, @@ -582,7 +587,7 @@ class BaseDecisionsConfig(ABC): ) def get_complete_url(self, api_base: str, model: str) -> str: - return f"{api_base.rstrip('/').removesuffix('/v1')}{self.path}" + return join_api_base(api_base.rstrip("/").removesuffix("/v1"), self.path) def signs_with_aws(self, api_base: str) -> bool: return False diff --git a/litellm/llms/base_llm/evals/transformation.py b/litellm/llms/base_llm/evals/transformation.py index 49124e271a6..40093953144 100644 --- a/litellm/llms/base_llm/evals/transformation.py +++ b/litellm/llms/base_llm/evals/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -79,7 +80,7 @@ class BaseEvalsAPIConfig(ABC): """ if api_base is None: raise ValueError("api_base is required") - return f"{api_base}/v1/{endpoint}" + return join_api_base(api_base, f"v1/{endpoint}") @abstractmethod def transform_create_eval_request( diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 5fe81a7f1b7..6399ded9555 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -121,7 +121,7 @@ class BaseTranslation(ABC): @staticmethod def merge_user_api_key_metadata_into_request( - request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place + request_data: dict[str, object], # mutable-ok: proxy hooks share and mutate the request payload dict in place user_api_key_dict: Optional["UserAPIKeyAuth"], ) -> None: """ diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 55a4b3da0b3..725146fe894 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -9,6 +9,8 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage +_DECODED_JSON: Final = TypeAdapter(object) + def _anthropic_stream_chunk_events(item: object) -> list[dict]: if isinstance(item, dict): @@ -30,7 +32,7 @@ def _anthropic_stream_chunk_events(item: object) -> list[dict]: if not payload or payload == "[DONE]": continue try: - parsed = json.loads(payload) + parsed = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError: continue if isinstance(parsed, dict): diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index 4faf0aaaf30..35b4bb90025 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -1,7 +1,7 @@ import types from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final import httpx from httpx._types import RequestFiles @@ -32,9 +32,10 @@ class BaseImageEditConfig(ABC): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/base_llm/interactions/transformation.py b/litellm/llms/base_llm/interactions/transformation.py index b7bbf7bde73..39ca1d8f7d7 100644 --- a/litellm/llms/base_llm/interactions/transformation.py +++ b/litellm/llms/base_llm/interactions/transformation.py @@ -11,7 +11,8 @@ Per OpenAPI spec (https://ai.google.dev/static/api/interactions.openapi.json): import types from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final import httpx @@ -60,9 +61,10 @@ class BaseInteractionsAPIConfig(ABC): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index d3b05716452..2663a3c42e7 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -141,7 +141,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): self, llm_router: Router, model: str, - request_data: dict[str, Any], + request_data: dict[str, object], litellm_parent_otel_span: Span, ) -> ResourceObjectType: """ @@ -399,7 +399,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): async def create_resource_for_each_model( self, llm_router: Router, - request_data: dict[str, Any], + request_data: dict[str, object], target_model_names_list: list[str], litellm_parent_otel_span: Span, ) -> list[ResourceObjectType]: diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 97d402e7a1b..171b193f60c 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -4,7 +4,7 @@ Base OCR types shared by the Rust OCR route and Python consumers. import builtins from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import PrivateAttr @@ -45,7 +45,7 @@ class OCRPageImage(LiteLLMPydanticObjectBase): """Image extracted from OCR page.""" image_base64: str | None = None - bbox: dict[str, Any] | None = None + bbox: dict[str, builtins.object] | None = None model_config = {"extra": "allow"} @@ -80,7 +80,7 @@ class OCRResponse(LiteLLMPydanticObjectBase): pages: list[OCRPage] model: str - document_annotation: Any | None = None + document_annotation: builtins.object | None = None usage_info: OCRUsageInfo | None = None content: str | None = None tables: list[dict[str, builtins.object]] | None = None diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index ac5032c08f1..7e58cc6e71b 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) class PassthroughMetadata(LiteLLMBaseModel): @@ -71,7 +72,7 @@ def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None def relayed_body(httpx_response: Response) -> str | dict: try: - body: Final[object] = httpx_response.json() + body: Final = _DECODED_JSON.validate_python(httpx_response.json()) except ValueError: return httpx_response.text return body if isinstance(body, dict) else httpx_response.text diff --git a/litellm/llms/base_llm/realtime/http_transformation.py b/litellm/llms/base_llm/realtime/http_transformation.py index 43a80edb493..a02054ec8cb 100644 --- a/litellm/llms/base_llm/realtime/http_transformation.py +++ b/litellm/llms/base_llm/realtime/http_transformation.py @@ -7,10 +7,11 @@ These are HTTP (not WebSocket) endpoints used by the WebRTC flow: """ from abc import ABC, abstractmethod -from typing import Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base + class BaseRealtimeHTTPConfig(ABC): """ @@ -61,7 +62,7 @@ class BaseRealtimeHTTPConfig(ABC): """Return the full URL for POST /realtime/transcription_sessions.""" base = (api_base or "").rstrip("/") base = base.removesuffix("/v1") - return f"{base}/v1/realtime/transcription_sessions" + return join_api_base(base, "v1/realtime/transcription_sessions") @abstractmethod def validate_environment( @@ -83,8 +84,7 @@ class BaseRealtimeHTTPConfig(ABC): def get_realtime_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: """Return the full URL for POST /realtime/calls (SDP exchange).""" - base: Final = (api_base or "").rstrip("/") - return f"{base}/v1/realtime/calls" + return join_api_base(api_base or "", "v1/realtime/calls") def get_realtime_calls_headers(self, ephemeral_key: str) -> dict: """ diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 5d2f92b5e82..7f246bf33da 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -80,7 +80,7 @@ class BaseRerankConfig(ABC): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py index 877edc15ec7..fe0e3abb87e 100644 --- a/litellm/llms/base_llm/sandbox/transformation.py +++ b/litellm/llms/base_llm/sandbox/transformation.py @@ -7,7 +7,7 @@ run code -> delete container; `code_interpreter_tool` combines all three. """ import builtins -from typing import Any, Final +from typing import Final import httpx from pydantic import Field, PrivateAttr @@ -42,8 +42,8 @@ class CodeExecutionResult(LiteLLMPydanticObjectBase): stdout: str = "" stderr: str = "" - results: list[dict[str, Any]] = Field(default_factory=list) - error: dict[str, Any] | None = None + results: list[dict[str, object]] = Field(default_factory=list) + error: dict[str, object] | None = None execution_count: int | None = None object: str = "code_execution" diff --git a/litellm/llms/base_llm/skills/transformation.py b/litellm/llms/base_llm/skills/transformation.py index ba202e3db5d..64a8e055832 100644 --- a/litellm/llms/base_llm/skills/transformation.py +++ b/litellm/llms/base_llm/skills/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic_skills import ( CreateSkillRequest, @@ -71,7 +72,7 @@ class BaseSkillsAPIConfig(ABC): """ if api_base is None: raise ValueError("api_base is required") - return f"{api_base}/v1/{endpoint}" + return join_api_base(api_base, f"v1/{endpoint}") @abstractmethod def transform_create_skill_request( diff --git a/litellm/llms/base_llm/text_to_speech/transformation.py b/litellm/llms/base_llm/text_to_speech/transformation.py index fb85cdb7687..be49098b62b 100644 --- a/litellm/llms/base_llm/text_to_speech/transformation.py +++ b/litellm/llms/base_llm/text_to_speech/transformation.py @@ -1,6 +1,7 @@ import types from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, TypedDict +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, TypedDict import httpx @@ -40,9 +41,10 @@ class BaseTextToSpeechConfig(ABC): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 28f6348e4d9..6885183aeba 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -152,7 +152,7 @@ class BaseVectorStoreConfig: api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: pass @@ -164,7 +164,7 @@ class BaseVectorStoreConfig: api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Optional async version of transform_search_vector_store_request. @@ -221,6 +221,11 @@ class BaseVectorStoreConfig: raise ValueError("api_base is required") return api_base + def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> str: + """Url of one vector store for retrieve, update and delete. Providers whose + ``get_complete_url`` adds a query string, like Azure, override this.""" + return f"{api_base}/{encoded_vector_store_id}" + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: from ..chat.transformation import BaseLLMException diff --git a/litellm/llms/baseten/chat.py b/litellm/llms/baseten/chat.py index 381d487c36d..c165831b90a 100644 --- a/litellm/llms/baseten/chat.py +++ b/litellm/llms/baseten/chat.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import litellm @@ -42,7 +43,7 @@ class BasetenConfig(OpenAIGPTConfig): frequency_penalty: int | None = None, stream_options: dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 48e79d39b8e..f37abacc883 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -174,12 +174,13 @@ class BedrockBatchesHandler: ) def job_status() -> "LiteLLMBatch": + auth_kwargs: Final[Mapping[str, object]] = dict(auth_params.model_dump()) return BedrockBatchesHandler.handle_model_invocation_job_status( batch_id=batch_id, aws_region_name=region, logging_obj=logging_obj, aws_bedrock_control_plane_endpoint=control_plane_endpoint, - **auth_params.model_dump(), + **auth_kwargs, ) try: diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 48b3795e274..18212b7390f 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, Union from urllib.parse import quote import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -18,6 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.a2a.common_utils import extract_text_from_a2a_response from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -50,6 +52,8 @@ else: HTTPHandler = Any AsyncHTTPHandler = Any +_DECODED_JSON: Final = TypeAdapter(object) + class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): def __init__(self, **kwargs): @@ -98,14 +102,14 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # Build the base endpoint URL for AgentCore # Note: We don't use get_runtime_endpoint as AgentCore has its own endpoint structure if aws_bedrock_runtime_endpoint: - base_url = aws_bedrock_runtime_endpoint + base_url = f"{aws_bedrock_runtime_endpoint}" else: base_url = f"https://bedrock-agentcore.{region}.{get_aws_dns_suffix(region)}" # Based on boto3 client.invoke_agent_runtime, the path is: # /runtimes/{URL-ENCODED-ARN}/invocations?qualifier= encoded_arn: Final = quote(agent_runtime_arn, safe="") - endpoint_url = f"{base_url}/runtimes/{encoded_arn}/invocations" + endpoint_url = join_api_base(base_url, f"runtimes/{encoded_arn}/invocations") # Add qualifier as query parameter if provided if "qualifier" in optional_params: @@ -313,7 +317,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return None try: - data: Final[object] = json.loads(json_str) + data: Final = _DECODED_JSON.validate_python(json.loads(json_str)) # Skip non-dict data (some lines contain JSON strings) return data if isinstance(data, dict) else None except json.JSONDecodeError: @@ -481,7 +485,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): if "application/json" in content_type: # Direct JSON response verbose_logger.debug("Parsing JSON response") - response_json: Final[object] = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) verbose_logger.debug("Response JSON: %s", response_json) return self._parse_json_response(response_json) else: @@ -565,7 +569,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): continue try: - data_obj = json.loads(json_str) + data_obj = _DECODED_JSON.validate_python(json.loads(json_str)) if not isinstance(data_obj, dict): continue @@ -705,7 +709,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) try: body: Final = response.read() - response_json: Final[object] = json.loads(body) + response_json: Final = _DECODED_JSON.validate_python(json.loads(body)) except (json.JSONDecodeError, Exception) as e: raise BedrockError( status_code=response.status_code, @@ -787,7 +791,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): continue try: - data_obj = json.loads(json_str) + data_obj = _DECODED_JSON.validate_python(json.loads(json_str)) if not isinstance(data_obj, dict): continue @@ -927,7 +931,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) try: body: Final = await response.aread() - response_json: Final[object] = json.loads(body) + response_json: Final = _DECODED_JSON.validate_python(json.loads(body)) except (json.JSONDecodeError, Exception) as e: raise BedrockError( status_code=response.status_code, diff --git a/litellm/llms/bedrock/chat/chat_completions/transformation.py b/litellm/llms/bedrock/chat/chat_completions/transformation.py index 907f8e50ac7..123e4a9e6b4 100644 --- a/litellm/llms/bedrock/chat/chat_completions/transformation.py +++ b/litellm/llms/bedrock/chat/chat_completions/transformation.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( inline_remote_image_urls, inline_remote_media, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token from litellm.llms.bedrock.common_utils import ( @@ -371,12 +372,9 @@ class AmazonBedrockRuntimeChatCompletionsConfig(OpenAILikeChatConfig): ), aws_region_name=aws_region_name, ) - base: Final = proxy_endpoint_url.rstrip("/") - if base.endswith("/openai/v1/chat/completions"): - return base - if base.endswith("/openai/v1"): - return f"{base}/chat/completions" - return f"{base}/openai/v1/chat/completions" + if proxy_endpoint_url.rstrip("/").endswith("/openai/v1"): + return join_api_base(proxy_endpoint_url, "chat/completions") + return ensure_api_route(proxy_endpoint_url, "openai/v1/chat/completions") def _params_with_region_from_path( self, optional_params: dict, model: str | None diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7d02b9ef230..a6f9192339f 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -10,6 +10,7 @@ from litellm.anthropic_beta_headers_manager import ( ) from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm import base_model_iterator from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -363,11 +364,11 @@ class BedrockConverseLLM(BaseAWSLLM): aws_region_name=aws_region_name, ) if (stream is not None and stream is True) and not fake_stream: - endpoint_url = f"{endpoint_url}/model/{modelId}/converse-stream" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse-stream" + endpoint_url = join_api_base(endpoint_url, f"model/{modelId}/converse-stream") + proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/converse-stream") else: - endpoint_url = f"{endpoint_url}/model/{modelId}/converse" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse" + endpoint_url = join_api_base(endpoint_url, f"model/{modelId}/converse") + proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/converse") ## COMPLETION CALL headers = {"Content-Type": "application/json"} diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index cd3727401ef..4d7ec663d09 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -12,7 +12,8 @@ from itertools import chain from typing import TYPE_CHECKING, Final, Literal, cast, overload import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -21,6 +22,7 @@ from litellm.constants import ( RESPONSE_FORMAT_TOOL_NAME, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, filter_exceptions_from_params, filter_internal_params, map_finish_reason, @@ -160,6 +162,14 @@ _TOOLS_AS_SENT: Final = TypeAdapter(tuple[Mapping[str, object], ...]) _JSON_VALUE: Final = TypeAdapter(object) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) _TOOLS_AS_RECEIVED: Final = TypeAdapter(list[object]) + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _ToolCallInContent(TypedDict): + name: ReadOnly[str] + + +_TOOL_CALL_IN_CONTENT: Final = TypeAdapter(_ToolCallInContent) CLAUDE_ON_BEDROCK_INVOKE_REMEDY: Final = " This model runs tool search on Bedrock through `bedrock/invoke/`." _DEPLOYMENT_PARAMS: Final = TypeAdapter(Mapping[str, object]) @@ -2453,9 +2463,17 @@ class AmazonConverseConfig(BaseConfig): if message.content is not None: try: tool_call_names: Final = self.get_tool_call_names(tools) - json_content: Final = json.loads(message.content) - if json_content.get("type") == "function" and json_content.get("name") in tool_call_names: - tool_calls: Final = [ChatCompletionMessageToolCall(function=Function(**json_content))] + json_content: Final = as_str_mapping(_JSON_VALUE.validate_python(json.loads(message.content))) + if ( + json_content is not None + and json_content.get("type") == "function" + and json_content.get("name") in tool_call_names + ): + tool_calls: Final = [ + ChatCompletionMessageToolCall( + function=Function(**_TOOL_CALL_IN_CONTENT.validate_python(json_content)) + ) + ] message.tool_calls = tool_calls message.content = None diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index 6877af74494..baef89912fa 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -15,7 +15,7 @@ from litellm._uuid import uuid from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError @@ -101,9 +101,10 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): encoded_agent_alias_id: Final = encode_url_path_segment(agent_alias_id, field_name="agent_alias_id") encoded_session_id: Final = encode_url_path_segment(session_id, field_name="session_id") - endpoint_url = f"{endpoint_url}/agents/{encoded_agent_id}/agentAliases/{encoded_agent_alias_id}/sessions/{encoded_session_id}/text" - - return endpoint_url + return join_api_base( + endpoint_url, + f"agents/{encoded_agent_id}/agentAliases/{encoded_agent_alias_id}/sessions/{encoded_session_id}/text", + ) def sign_request( self, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_ai21_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_ai21_transformation.py index 3005bc82299..777237b4f18 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_ai21_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_ai21_transformation.py @@ -56,9 +56,10 @@ class AmazonAI21Config(AmazonInvokeConfig, BaseConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py index 00c2eaf0f8f..2118d183b22 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py @@ -37,9 +37,10 @@ class AmazonCohereConfig(AmazonInvokeConfig, CohereChatConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_llama_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_llama_transformation.py index d2eb6eef9ef..cb867a09f8f 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_llama_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_llama_transformation.py @@ -37,9 +37,10 @@ class AmazonLlamaConfig(AmazonInvokeConfig, BaseConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py index 44511eeaa3e..4420bb3a418 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py @@ -47,9 +47,10 @@ class AmazonMistralConfig(AmazonInvokeConfig, BaseConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index 4bf1a1cba73..f1e1f3e8738 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.request_metadata import ( @@ -97,13 +98,8 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): # Encode model ID for ARNs (e.g., :imported-model/ -> :imported-model%2F) model_id = CommonUtils.encode_bedrock_runtime_modelid_arn(model_id) - # Build the invoke URL - if stream: - endpoint_url = f"{endpoint_url}/model/{model_id}/invoke-with-response-stream" - else: - endpoint_url = f"{endpoint_url}/model/{model_id}/invoke" - - return endpoint_url + invoke_route: Final = "invoke-with-response-stream" if stream else "invoke" + return join_api_base(endpoint_url, f"model/{model_id}/{invoke_route}") def sign_request( self, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py index a18e97c63a6..736b3ea3cd2 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py @@ -43,9 +43,10 @@ class AmazonTitanConfig(AmazonInvokeConfig, BaseConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py index b29ddf01d1e..2655c13be44 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py @@ -44,9 +44,10 @@ class AmazonAnthropicConfig(AmazonInvokeConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 87d4fc0b32e..51b5a99272e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( prompt_factory, render_custom_prompt_template, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call from litellm.llms.bedrock.common_utils import BedrockError @@ -136,20 +137,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): aws_bedrock_runtime_endpoint: Final = optional_params.get( "aws_bedrock_runtime_endpoint", None ) # https://bedrock-runtime.{region_name}.amazonaws.com - endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( + endpoint_url, _ = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_region_name=self._get_aws_region_name(optional_params=optional_params, model=model), ) - - if (stream is not None and stream is True) and provider != "ai21": - endpoint_url = f"{endpoint_url}/model/{modelId}/invoke-with-response-stream" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke-with-response-stream" - else: - endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" - - return endpoint_url + invoke_route: Final = "invoke-with-response-stream" if stream is True and provider != "ai21" else "invoke" + return join_api_base(endpoint_url, f"model/{modelId}/{invoke_route}") def sign_request( self, diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index e48196946cd..a1aa2034e4c 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -21,6 +21,7 @@ from litellm.types.utils import ModelResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.anthropic.chat.handler import ModelResponseIterator LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -134,7 +135,7 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "ModelResponseIterator": return self.get_model_response_iterator_for_request( streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, optional_params={} ) @@ -145,7 +146,8 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): sync_stream: bool, json_mode: bool | None, optional_params: Mapping[str, object], - ) -> object: + logging_obj: LiteLLMLoggingObj | None = None, + ) -> "ModelResponseIterator": from litellm.llms.anthropic.chat.handler import ModelResponseIterator return ModelResponseIterator( diff --git a/litellm/llms/bedrock/claude_platform/common_utils.py b/litellm/llms/bedrock/claude_platform/common_utils.py index e25e58056f5..77cae177792 100644 --- a/litellm/llms/bedrock/claude_platform/common_utils.py +++ b/litellm/llms/bedrock/claude_platform/common_utils.py @@ -5,6 +5,7 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError from litellm.secret_managers.main import get_secret_str @@ -140,9 +141,7 @@ class BedrockClaudePlatformMixin(BaseAWSLLM): if api_base is None: aws_region_name: Final = self._get_required_aws_region_name(optional_params) api_base = f"https://{CLAUDE_PLATFORM_SERVICE_NAME}.{aws_region_name}.api.aws" - if not api_base.endswith("/v1/messages"): - api_base = f"{api_base.rstrip('/')}/v1/messages" - return api_base + return ensure_api_route(api_base, "v1/messages") def sign_request( self, diff --git a/litellm/llms/bedrock/claude_platform/transformation.py b/litellm/llms/bedrock/claude_platform/transformation.py index cea1140bf65..0194f1c7e45 100644 --- a/litellm/llms/bedrock/claude_platform/transformation.py +++ b/litellm/llms/bedrock/claude_platform/transformation.py @@ -14,6 +14,7 @@ from .common_utils import ( ) if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.anthropic.chat.handler import ModelResponseIterator @@ -110,6 +111,7 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig): sync_stream: bool, json_mode: bool | None, optional_params: Mapping[str, object], + logging_obj: "Logging | None" = None, ) -> "ModelResponseIterator": from litellm.llms.anthropic.chat.handler import ModelResponseIterator diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 0cf3e895ad2..2b946d9664c 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -31,6 +31,7 @@ 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.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -1305,10 +1306,7 @@ def build_mantle_messages_url( """ override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env() if override: - base: Final = override.rstrip("/") - if base.endswith(MANTLE_MESSAGES_PATH): - return base - return f"{base}{MANTLE_MESSAGES_PATH}" + return ensure_api_route(override, MANTLE_MESSAGES_PATH) return f"https://bedrock-mantle.{region}.api.aws{MANTLE_MESSAGES_PATH}" @@ -2023,8 +2021,13 @@ def get_bedrock_response_stream_shape(): _BEDROCK_STREAM_OUTPUT_SHAPES: Final = ("ConverseStreamOutput", "ResponseStream") +class _ModeledErrorMetadata(TypedDict, total=False): + httpStatusCode: ReadOnly[int] + + def _modeled_error_status(member: Shape) -> int | None: - status: Final = (member.metadata or {}).get("error", {}).get("httpStatusCode") + error: Final[_ModeledErrorMetadata] = (member.metadata or {}).get("error", {}) + status: Final = error.get("httpStatusCode") return None if status is None else int(status) diff --git a/litellm/llms/bedrock/cost_calculation.py b/litellm/llms/bedrock/cost_calculation.py index 0a2e39ab972..1b45531dda4 100644 --- a/litellm/llms/bedrock/cost_calculation.py +++ b/litellm/llms/bedrock/cost_calculation.py @@ -3,9 +3,13 @@ Helper util for handling bedrock-specific cost calculation - e.g.: prompt caching """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + generic_cost_per_token, + get_provider_specific_geo_multiplier, +) if TYPE_CHECKING: from litellm.types.utils import Usage @@ -17,9 +21,13 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None) Follows the same logic as Anthropic's cost per token calculation. """ - return generic_cost_per_token( + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="bedrock") + prompt_cost, completion_cost = generic_cost_per_token( model=model, usage=usage, custom_llm_provider="bedrock", service_tier=service_tier, + model_info=model_info, ) + geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) + return prompt_cost * geo_multiplier, completion_cost * geo_multiplier diff --git a/litellm/llms/bedrock/count_tokens/kimi_k3.py b/litellm/llms/bedrock/count_tokens/kimi_k3.py index 5e3fd56e4ba..4f8b4624cd8 100644 --- a/litellm/llms/bedrock/count_tokens/kimi_k3.py +++ b/litellm/llms/bedrock/count_tokens/kimi_k3.py @@ -27,6 +27,7 @@ from typing_extensions import Self import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.litellm_core_utils.prompt_templates.factory import ( chat_completion_tools_for, is_flat_responses_tool, @@ -67,12 +68,14 @@ class KimiK3TokenizerSource: def _hugging_face_ranks_path() -> str: import huggingface_hub - return huggingface_hub.hf_hub_download( - repo_id=KIMI_K3_REPO, - filename=KIMI_K3_RANKS_FILE, - revision=KIMI_K3_REVISION, - local_files_only=litellm.disable_hf_tokenizer_download is True, - ) + if litellm.disable_hf_tokenizer_download is True or is_offline_mode(): + cached: Final = huggingface_hub.try_to_load_from_cache( + repo_id=KIMI_K3_REPO, filename=KIMI_K3_RANKS_FILE, revision=KIMI_K3_REVISION + ) + if not isinstance(cached, str): + raise FileNotFoundError(f"{KIMI_K3_RANKS_FILE} is not in the Hugging Face cache") + return cached + return huggingface_hub.hf_hub_download(repo_id=KIMI_K3_REPO, filename=KIMI_K3_RANKS_FILE, revision=KIMI_K3_REVISION) HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256) diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 14edd39c386..34e622b5e8d 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -7,6 +7,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url @@ -43,7 +44,7 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): messages_url: Final = build_mantle_messages_url( api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name ) - return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}" + return join_api_base(messages_url, MANTLE_COUNT_TOKENS_SUFFIX) async def handle_count_tokens_request( self, diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index 7cca4f69acc..a9f95c1f9b9 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -11,6 +11,7 @@ from typing import Final, Literal from pydantic import JsonValue +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model @@ -236,9 +237,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_region_name=aws_region_name, ) - endpoint: Final = f"{base_url}/model/{encoded_model_id}/count-tokens" - - return endpoint + return join_api_base(base_url, f"model/{encoded_model_id}/count-tokens") def transform_bedrock_response_to_anthropic( self, bedrock_response: Mapping[str, JsonValue] diff --git a/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py index 581b89369f7..fba2673b5aa 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py @@ -35,9 +35,10 @@ class AmazonTitanG1Config: @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py index 3b594662a9b..7d319ca6b06 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py @@ -39,9 +39,10 @@ class AmazonTitanV2Config: @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 58cfb6cb95e..4d27d94664f 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -12,6 +12,7 @@ import httpx import litellm from litellm.constants import BEDROCK_EMBEDDING_PROVIDERS_LITERAL +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.cohere.embed.handler import embedding as cohere_embedding from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -512,9 +513,9 @@ class BedrockEmbedding(BaseAWSLLM): aws_region_name=aws_region_name, ) if has_async_invoke: - endpoint_url = f"{endpoint_url}/async-invoke" + endpoint_url = join_api_base(endpoint_url, "async-invoke") else: - endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" + endpoint_url = join_api_base(endpoint_url, f"model/{modelId}/invoke") if batch_data is not None: if aembedding: @@ -611,7 +612,7 @@ class BedrockEmbedding(BaseAWSLLM): # Encode the ARN for use in URL path encoded_arn: Final = quote(invocation_arn, safe="") - status_url: Final = f"{endpoint_url.rstrip('/')}/async-invoke/{encoded_arn}" + status_url: Final = join_api_base(endpoint_url, f"async-invoke/{encoded_arn}") # Prepare headers for GET request headers: Final = {"Content-Type": "application/json"} diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 3d23b69f846..37f95ff7b9b 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -1,6 +1,6 @@ import asyncio from collections.abc import Coroutine, Mapping -from typing import Any, Final +from typing import Final import httpx @@ -139,7 +139,7 @@ class BedrockFilesHandler(BaseAWSLLM): optional_params: dict, timeout: float | httpx.Timeout, max_retries: int | None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: """ Download file content from S3 bucket for Bedrock files. Supports both sync and async operations. diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 53bb332f511..bdfe6f5a8d2 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -10,10 +10,12 @@ import json from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.image_edit.amazon_nova_canvas_image_edit_transformation import ( BedrockAmazonNovaCanvasImageEditConfig, ) @@ -38,6 +40,8 @@ if TYPE_CHECKING: else: AWSPreparedRequest = Any +_DECODED_JSON: Final = TypeAdapter(object) + class BedrockImageEditPreparedRequest(LiteLLMBaseModel): """ @@ -226,7 +230,7 @@ class BedrockImageEdit(BaseAWSLLM): aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint, aws_region_name=boto3_credentials_info.aws_region_name, ) - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" + proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/invoke") data: Final = self._get_request_body( model=model, image=image, @@ -316,7 +320,7 @@ class BedrockImageEdit(BaseAWSLLM): additional_args={"complete_input_dict": data}, ) verbose_logger.debug("raw model_response: %s", response.text) - response_dict: Final = response.json() + response_dict: Final = _DECODED_JSON.validate_python(response.json()) if response_dict is None: raise ValueError("Error in response object format, got None") diff --git a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py index acc847e0e3c..e5c5edba750 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py @@ -1,4 +1,5 @@ import types +from collections.abc import Mapping from typing import Final from openai.types.image import Image @@ -20,9 +21,10 @@ class AmazonStability3Config: @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py index 145e6a0781f..5731541da56 100644 --- a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py @@ -3,6 +3,7 @@ Transformation logic for Amazon Titan Image Generation. """ import types +from collections.abc import Mapping from typing import Final from openai.types.image import Image @@ -35,16 +36,17 @@ class AmazonTitanImageGenerationConfig: width: int | None = None, height: int | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index e81dfe76942..abd6d40f9e9 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -8,6 +8,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.image_generation.amazon_nova_canvas_transformation import ( AmazonNovaCanvasConfig, ) @@ -248,7 +249,7 @@ class BedrockImageGeneration(BaseAWSLLM): aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint, aws_region_name=boto3_credentials_info.aws_region_name, ) - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" + proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/invoke") data: Final = self._get_request_body( model=model, prompt=prompt, 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 12421aa94fd..805652630ed 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -1,4 +1,4 @@ -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast @@ -145,7 +145,7 @@ class AmazonAnthropicClaudeMessagesConfig( self, headers: dict, model: str, - messages: list[object], + messages: Sequence[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index c71721d800e..d4a2d2eabe9 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -7,7 +7,7 @@ stripping, and the anthropic-version / anthropic-beta placement (headers, never the body) that are specific to the bedrock-mantle endpoint. """ -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -90,7 +90,7 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/bedrock/model_listing.py b/litellm/llms/bedrock/model_listing.py index 99ff48f9400..409bf2c7b97 100644 --- a/litellm/llms/bedrock/model_listing.py +++ b/litellm/llms/bedrock/model_listing.py @@ -4,7 +4,7 @@ from typing import Final, TypeAlias from urllib.parse import quote, urlencode import httpx -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token @@ -15,6 +15,7 @@ from litellm.types.llms.bedrock import AwsAuthParams LIST_MODELS_TIMEOUT: Final = 10.0 INFERENCE_PROFILES_PAGE_SIZE: Final = 1000 INFERENCE_PROFILES_PAGE_CAP: Final = 20 +_DECODED_JSON: Final = TypeAdapter(object) QueryPairs: TypeAlias = tuple[tuple[str, str], ...] @@ -103,7 +104,7 @@ class BedrockModelLister(BaseAWSLLM): f"Status code: {response.status_code}, Response: {response.text}" ), ) - return response.json() + return _DECODED_JSON.validate_python(response.json()) def _auth_headers(self, url: str) -> Mapping[str, str]: if self._bearer_token is not None: diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index a90b4829634..1b01b0ec6d0 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -544,7 +544,7 @@ class _StreamFrame(TypedDict): def _unpack_uint32(buffer: bytes) -> int: - return struct.unpack("!I", buffer)[0] + return int.from_bytes(buffer, "big") def _collect_stream_delta_text_holders(delta: object) -> list[_DeltaHolder]: @@ -1255,7 +1255,9 @@ 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 = _BODY_BY_REFERENCE.validate_python( + 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")) ): diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 01ef5af2b67..0634275135b 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -6,6 +6,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -160,7 +161,7 @@ class BedrockRerankHandler(BaseAWSLLM): aws_region_name=boto3_credentials_info.aws_region_name, ) proxy_endpoint_url = proxy_endpoint_url.replace("bedrock-runtime", "bedrock-agent-runtime") - proxy_endpoint_url = f"{proxy_endpoint_url}/rerank" + proxy_endpoint_url = join_api_base(proxy_endpoint_url, "rerank") body: Final = json.dumps(data).encode("utf-8") headers = {"Content-Type": "application/json"} diff --git a/litellm/llms/bedrock/responses/transformation.py b/litellm/llms/bedrock/responses/transformation.py index 086d1211835..7cdd7b08f99 100644 --- a/litellm/llms/bedrock/responses/transformation.py +++ b/litellm/llms/bedrock/responses/transformation.py @@ -46,6 +46,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, convert_url_to_base64, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -223,9 +224,10 @@ class BedrockOpenAIResponsesConfig(BaseAWSLLM, OpenAIResponsesAPIConfig): litellm_params: dict, # mutable-ok: signature fixed by the BaseResponsesAPIConfig override contract ) -> str: region: Final = self._get_aws_region_name(optional_params=litellm_params, model=None) + configured_endpoint: Final = litellm_params.get("aws_bedrock_runtime_endpoint") override: Final = ( api_base - or litellm_params.get("aws_bedrock_runtime_endpoint") + or (f"{configured_endpoint}" if configured_endpoint else None) or get_secret_str("AWS_BEDROCK_RUNTIME_ENDPOINT") ) # Partition-aware: bedrock-runtime is amazonaws.com.cn in China, and other @@ -237,7 +239,7 @@ class BedrockOpenAIResponsesConfig(BaseAWSLLM, OpenAIResponsesAPIConfig): (host[: -len(suffix)] for suffix in BEDROCK_RUNTIME_OPENAI_BASE_SUFFIXES if host.endswith(suffix)), host, ) - return f"{base}{BEDROCK_RUNTIME_OPENAI_RESPONSES_PATH}" + return join_api_base(base, BEDROCK_RUNTIME_OPENAI_RESPONSES_PATH) def supports_native_file_search(self) -> bool: return False diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index ba8ce7e5625..96d1da19bcc 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -6,7 +6,7 @@ from urllib.parse import urlparse import httpx from litellm._logging import verbose_logger -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError @@ -196,7 +196,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): aws_region_name=self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=aws_region_name), endpoint_type="agent", ) - return f"{endpoint_url}/knowledgebases" + return join_api_base(endpoint_url, "knowledgebases") def transform_search_vector_store_request( self, @@ -212,7 +212,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): query = " ".join(query) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}/retrieve" + url: Final = join_api_base(api_base, f"{encoded_vector_store_id}/retrieve") request_body: Final[dict[str, object]] = { "retrievalQuery": BedrockKBRetrievalQuery(text=query), diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py index 78881153e91..aa6510f0491 100644 --- a/litellm/llms/bedrock_mantle/messages/transformation.py +++ b/litellm/llms/bedrock_mantle/messages/transformation.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import MANTLE_MESSAGES_PATH from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig @@ -33,7 +34,7 @@ def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mappi configured, ) host: Final = f"https://bedrock-mantle.{region}.api.aws" if MANTLE_HOST_RE.match(stripped) else stripped - return f"{host}{MANTLE_MESSAGES_PATH}" + return join_api_base(host, MANTLE_MESSAGES_PATH) class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleMessagesConfig): diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 8a0ce58c8b2..3532cfbd963 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -22,6 +22,7 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items from litellm.llms.base_llm.responses.transformation import ReasoningSummaryLevel @@ -100,7 +101,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI if MANTLE_HOST_RE.match(base): base = f"https://bedrock-mantle.{region}.api.aws" path: Final = "/openai/v1/responses" if self.use_openai_path else "/v1/responses" - return f"{base}{path}" + return join_api_base(base, path) def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 66eacc97fba..38bcbb0d905 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -19,7 +19,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import async_safe_get, join_api_base, safe_get from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -51,6 +51,7 @@ class _BflPolledResult(TypedDict): _BFL_POLLED_RESULT_ADAPTER: Final = TypeAdapter(_BflPolledResult) +_DECODED_JSON: Final = TypeAdapter(object) _BFL_REQUEST_PARAMS: Final = ( @@ -189,11 +190,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): """ Get the complete URL for the Black Forest Labs API request. """ - base_url: str = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE endpoint: Final = self._get_model_endpoint(model) - return f"{base_url}{endpoint}" + return join_api_base(base_url, endpoint) def _read_image_bytes( self, @@ -317,7 +317,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): The response contains: {"status": "Ready", "result": {"sample": "https://..."}} """ try: - raw_response_data: Final[object] = raw_response.json() + raw_response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise BlackForestLabsError( status_code=raw_response.status_code, diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index 9d586fba26a..52f2c9a7c50 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -39,6 +40,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): @@ -202,11 +204,10 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the Black Forest Labs API request. """ - base_url: str = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE endpoint: Final = self._get_model_endpoint(model) - return f"{base_url}{endpoint}" + return join_api_base(base_url, endpoint) def transform_image_generation_request( self, @@ -272,7 +273,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): The response contains: {"status": "Ready", "result": {"sample": "https://..."}} """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise BlackForestLabsError( status_code=raw_response.status_code, diff --git a/litellm/llms/cerebras/chat.py b/litellm/llms/cerebras/chat.py index c3aa26ade35..9fc99500fc5 100644 --- a/litellm/llms/cerebras/chat.py +++ b/litellm/llms/cerebras/chat.py @@ -4,6 +4,7 @@ Cerebras Chat Completions API this is OpenAI compatible - no translation needed / occurs """ +from collections.abc import Mapping from typing import Final from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig @@ -41,7 +42,7 @@ class CerebrasConfig(OpenAIGPTConfig): user: str | None = None, reasoning_effort: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index ff9c6476cba..a5fb1c9c6b9 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -1,9 +1,9 @@ -import base64 import json import os import time from collections.abc import Mapping -from typing import Final, TypeAlias +from functools import lru_cache +from typing import TYPE_CHECKING, Final import httpx from pydantic import JsonValue, TypeAdapter, ValidationError @@ -24,19 +24,21 @@ from .common_utils import ( CHATGPT_OAUTH_TOKEN_URL, GetAccessTokenError, GetDeviceCodeError, + JsonObject, RefreshAccessTokenError, + decode_jwt_claims, + extract_chatgpt_account_id, ) +if TYPE_CHECKING: + from litellm.types.router import GenericLiteLLMParams + TOKEN_EXPIRY_SKEW_SECONDS: Final = 60 TOKEN_REFRESH_TIMEOUT_SECONDS: Final = 30 DEVICE_CODE_TIMEOUT_SECONDS: Final = 15 * 60 DEVICE_CODE_COOLDOWN_SECONDS: Final = 5 * 60 DEVICE_CODE_POLL_SLEEP_SECONDS: Final = 5 -OPENAI_AUTH_CLAIM_KEY: Final = "https://api.openai.com/auth" - -JsonObject: TypeAlias = Mapping[str, JsonValue] - _JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject) @@ -44,6 +46,26 @@ def _optional_str(value: JsonValue | None) -> str | None: return value if isinstance(value, str) else None +def get_chatgpt_auth_file( + litellm_params: "Mapping[str, object] | GenericLiteLLMParams | None", +) -> str | None: + if litellm_params is None: + return None + value: Final = ( + litellm_params.get("chatgpt_auth_file") + if isinstance(litellm_params, Mapping) + else litellm_params.chatgpt_auth_file + ) + if not isinstance(value, str) or not value: + return None + return os.path.abspath(os.path.expanduser(value)) + + +@lru_cache(maxsize=128) +def get_cached_authenticator(auth_file: str) -> "Authenticator": + return Authenticator(auth_file=auth_file) + + def _token_refresh_timeout() -> httpx.Timeout: configured: Final = get_configured_request_timeout() seconds: Final = ( @@ -53,13 +75,13 @@ def _token_refresh_timeout() -> httpx.Timeout: class Authenticator: - def __init__(self) -> None: - self.token_dir = os.getenv( - "CHATGPT_TOKEN_DIR", - os.path.expanduser("~/.config/litellm/chatgpt"), + def __init__(self, auth_file: str | None = None) -> None: + default_auth_file: Final = os.path.join( + os.getenv("CHATGPT_TOKEN_DIR", os.path.expanduser("~/.config/litellm/chatgpt")), + os.getenv("CHATGPT_AUTH_FILE", "auth.json"), ) - self.auth_file = os.path.join(self.token_dir, os.getenv("CHATGPT_AUTH_FILE", "auth.json")) - self._ensure_token_dir() + self.auth_file = os.path.abspath(os.path.expanduser(auth_file or default_auth_file)) + self.token_dir = os.path.dirname(self.auth_file) def get_api_base(self) -> str: return os.getenv("CHATGPT_API_BASE") or os.getenv("OPENAI_CHATGPT_API_BASE") or CHATGPT_API_BASE @@ -99,16 +121,18 @@ class Authenticator: tokens: Final = self._login_device_code() return tokens["access_token"] - def get_account_id(self) -> str | None: + def get_account_id(self, access_token: str | None = None) -> str | None: auth_data: Final = self._read_auth_file() + if access_token and (not auth_data or auth_data.get("access_token") != access_token): + return extract_chatgpt_account_id(access_token) if not auth_data: return None account_id: Final = _optional_str(auth_data.get("account_id")) if account_id: return account_id id_token: Final = auth_data.get("id_token") - access_token: Final = auth_data.get("access_token") - derived: Final = self._extract_account_id(_optional_str(id_token or access_token)) + saved_access_token: Final = auth_data.get("access_token") + derived: Final = extract_chatgpt_account_id(_optional_str(id_token or saved_access_token)) if derived: self._write_auth_file({**auth_data, "account_id": derived}) return derived @@ -129,6 +153,7 @@ class Authenticator: def _write_auth_file(self, data: JsonObject) -> None: try: + self._ensure_token_dir() with open(self.auth_file, "w") as f: json.dump(data, f) except OSError as exc: @@ -146,34 +171,11 @@ class Authenticator: return time.time() >= float(derived_expires_at) - TOKEN_EXPIRY_SKEW_SECONDS def _get_expires_at(self, token: str) -> int | None: - claims: Final = self._decode_jwt_claims(token) - exp: Final = claims.get("exp") + exp: Final = decode_jwt_claims(token).get("exp") if isinstance(exp, (int, float)): return int(exp) return None - def _decode_jwt_claims(self, token: str) -> JsonObject: - try: - parts: Final = token.split(".") - if len(parts) < 2: - return {} - payload_b64: Final = parts[1] + "=" * (-len(parts[1]) % 4) - payload_bytes: Final = base64.urlsafe_b64decode(payload_b64) - return _JSON_OBJECT_ADAPTER.validate_python(json.loads(payload_bytes.decode("utf-8"))) - except Exception: - return {} - - def _extract_account_id(self, token: str | None) -> str | None: - if not token: - return None - claims: Final = self._decode_jwt_claims(token) - auth_claims: Final = claims.get(OPENAI_AUTH_CLAIM_KEY) - if isinstance(auth_claims, dict): - account_id: Final = auth_claims.get("chatgpt_account_id") - if isinstance(account_id, str) and account_id: - return account_id - return None - def _login_device_code(self) -> dict[str, str]: cooldown_remaining: Final = self._get_device_code_cooldown_remaining(self._read_auth_file()) if cooldown_remaining > 0: @@ -361,15 +363,16 @@ class Authenticator: "refresh_token": _optional_str(data.get("refresh_token")) or refresh_token, "id_token": id_token, } + account_id: Final = self.get_account_id() auth_data: Final = self._build_auth_record(refreshed) - self._write_auth_file(auth_data) + self._write_auth_file({**auth_data, "account_id": account_id} if account_id else auth_data) return refreshed def _build_auth_record(self, tokens: dict[str, str]) -> JsonObject: access_token: Final = tokens.get("access_token") id_token: Final = tokens.get("id_token") expires_at: Final = self._get_expires_at(access_token) if access_token else None - account_id: Final = self._extract_account_id(id_token or access_token) + account_id: Final = extract_chatgpt_account_id(id_token or access_token) return { "access_token": access_token, "refresh_token": tokens.get("refresh_token"), diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index ee3120bacea..49be583f6d1 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -7,6 +7,8 @@ Normalizes non-spec-compliant tool_call chunks from the ChatGPT backend API. from collections.abc import Awaitable from typing import Final, Protocol +from pydantic import TypeAdapter + from litellm.types.utils import ( ChatCompletionDeltaCustomToolCall, ChatCompletionDeltaToolCall, @@ -23,6 +25,9 @@ class ChatGPTChunkStream(Protocol): def __anext__(self) -> Awaitable[ModelResponseStream]: ... +_DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) + + def _first_choice_delta(chunk: ModelResponseStream) -> Delta | None: return chunk.choices[0].delta @@ -47,7 +52,7 @@ class ChatGPTToolCallNormalizer: self._last_id: str | None = None # tracks which tool call the next delta belongs to def __getattr__(self, name: str) -> object: - return getattr(self._stream, name) + return _DELEGATED_ATTRIBUTE.validate_python(getattr(self._stream, name)) def __iter__(self) -> "ChatGPTToolCallNormalizer": return self diff --git a/litellm/llms/chatgpt/chat/transformation.py b/litellm/llms/chatgpt/chat/transformation.py index a1aac7f315b..9005185931f 100644 --- a/litellm/llms/chatgpt/chat/transformation.py +++ b/litellm/llms/chatgpt/chat/transformation.py @@ -1,14 +1,20 @@ +from collections.abc import Mapping from typing import Any, Final from litellm.exceptions import AuthenticationError from litellm.llms.openai.openai import OpenAIConfig from litellm.types.llms.openai import AllMessageValues +from litellm.types.router import GenericLiteLLMParams -from ..authenticator import Authenticator +from ..authenticator import Authenticator, get_cached_authenticator, get_chatgpt_auth_file from ..common_utils import ( GetAccessTokenError, + chatgpt_client_api_base, ensure_chatgpt_session_id, + get_chatgpt_client_credential, get_chatgpt_default_headers, + get_chatgpt_request_credential, + strip_chatgpt_client_credential_headers, ) from .streaming_utils import ChatGPTToolCallNormalizer @@ -19,9 +25,27 @@ class ChatGPTConfig(OpenAIConfig): api_key: str | None = None, api_base: str | None = None, custom_llm_provider: str = "openai", + authenticator: Authenticator | None = None, ) -> None: super().__init__() - self.authenticator = Authenticator() + self.authenticator = authenticator if authenticator is not None else Authenticator() + + def _resolve_authenticator( + self, litellm_params: Mapping[str, object] | GenericLiteLLMParams | None + ) -> Authenticator: + auth_file: Final = get_chatgpt_auth_file(litellm_params) + return get_cached_authenticator(auth_file) if auth_file else self.authenticator + + @staticmethod + def _get_access_token_or_raise(authenticator: Authenticator, model: str, llm_provider: str) -> str: + try: + return authenticator.get_access_token() + except GetAccessTokenError as e: + raise AuthenticationError( + model=model, + llm_provider=llm_provider, + message=str(e), + ) def api_base_without_login(self) -> str: return self.authenticator.get_api_base() @@ -32,16 +56,16 @@ class ChatGPTConfig(OpenAIConfig): api_base: str | None, api_key: str | None, custom_llm_provider: str, + litellm_params: Mapping[str, object] | GenericLiteLLMParams | None = None, ) -> tuple[str | None, str | None, str]: - dynamic_api_base: Final = self.api_base_without_login() - try: - dynamic_api_key: Final = self.authenticator.get_access_token() - except GetAccessTokenError as e: - raise AuthenticationError( - model=model, - llm_provider=custom_llm_provider, - message=str(e), - ) + client_credential: Final = get_chatgpt_request_credential(litellm_params) + if client_credential is not None: + return chatgpt_client_api_base(api_base), client_credential.access_token, custom_llm_provider + if api_key: + return api_base or self.api_base_without_login(), api_key, custom_llm_provider + authenticator: Final = self._resolve_authenticator(litellm_params) + dynamic_api_base: Final = authenticator.get_api_base() + dynamic_api_key: Final = self._get_access_token_or_raise(authenticator, model, custom_llm_provider) return dynamic_api_base, dynamic_api_key, custom_llm_provider def get_openai_compatible_provider_info( @@ -50,26 +74,49 @@ class ChatGPTConfig(OpenAIConfig): api_base: str | None, api_key: str | None, custom_llm_provider: str, + litellm_params: Mapping[str, object] | GenericLiteLLMParams | None = None, ) -> tuple[str | None, str | None, str]: - return self._get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider) + return self._get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider, litellm_params) def validate_environment( self, - headers: dict, + headers: dict[str, str], model: str, messages: list[AllMessageValues], optional_params: dict, - litellm_params: dict, + litellm_params: dict[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: + client_credential: Final = get_chatgpt_client_credential(headers) + if client_credential is not None: + return dict( + get_chatgpt_default_headers( + client_credential.access_token, + client_credential.account_id, + ensure_chatgpt_session_id(litellm_params), + ), + **strip_chatgpt_client_credential_headers(headers), + ) + authenticator: Final = self._resolve_authenticator(litellm_params) + resolved_api_key: Final = ( + self._get_access_token_or_raise(authenticator, model, "chatgpt") + if get_chatgpt_auth_file(litellm_params) + else api_key + ) validated_headers: Final = super().validate_environment( - headers, model, messages, optional_params, litellm_params, api_key, api_base + dict(strip_chatgpt_client_credential_headers(headers)), + model, + messages, + optional_params, + litellm_params, + resolved_api_key, + api_base, ) - account_id: Final = self.authenticator.get_account_id() + account_id: Final = authenticator.get_account_id(resolved_api_key) session_id: Final = ensure_chatgpt_session_id(litellm_params) - default_headers: Final = get_chatgpt_default_headers(api_key or "", account_id, session_id) + default_headers: Final = get_chatgpt_default_headers(resolved_api_key or "", account_id, session_id) return {**default_headers, **validated_headers} def post_stream_processing(self, stream: Any) -> ChatGPTToolCallNormalizer: diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index fe33219f110..ec7fbfa92da 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -2,14 +2,32 @@ Constants and helpers for ChatGPT subscription OAuth. """ +import base64 +import json import os import platform -from typing import Any, Final +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias +from urllib.parse import urlsplit from uuid import uuid4 import httpx +from pydantic import JsonValue, TypeAdapter +from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.utils import ProviderSpecificHeader + +if TYPE_CHECKING: + from litellm.types.router import GenericLiteLLMParams + +_SCOPED_HEADERS: Final[TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None]] = TypeAdapter( + ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None +) +_STRING_HEADERS: Final = TypeAdapter(Mapping[str, str]) +_PARAMS: Final = TypeAdapter(Mapping[str, object]) # OAuth + API constants (derived from openai/codex) CHATGPT_AUTH_BASE: Final = "https://auth.openai.com" @@ -19,6 +37,12 @@ CHATGPT_OAUTH_TOKEN_URL: Final = f"{CHATGPT_AUTH_BASE}/oauth/token" CHATGPT_DEVICE_VERIFY_URL: Final = f"{CHATGPT_AUTH_BASE}/codex/device" CHATGPT_API_BASE: Final = "https://chatgpt.com/backend-api/codex" CHATGPT_CLIENT_ID: Final = "app_EMoamEEZ73f0CkXaXp7hrann" +CHATGPT_ACCOUNT_ID_HEADER: Final = "ChatGPT-Account-Id" +OPENAI_AUTH_CLAIM_KEY: Final = "https://api.openai.com/auth" + +JsonObject: TypeAlias = Mapping[str, JsonValue] +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject) +_EMPTY_CLAIMS: Final[JsonObject] = MappingProxyType({}) DEFAULT_ORIGINATOR: Final = "codex_cli_rs" DEFAULT_USER_AGENT: Final = "codex_cli_rs/0.0.0 (Unknown 0; unknown) unknown" @@ -229,7 +253,7 @@ def get_chatgpt_default_headers( access_token: str, account_id: str | None, session_id: str | None = None, -) -> dict: +) -> dict[str, str]: originator: Final = get_chatgpt_originator() user_agent: Final = get_chatgpt_user_agent(originator) headers: Final = { @@ -242,30 +266,133 @@ def get_chatgpt_default_headers( if session_id: headers["session_id"] = session_id if account_id: - headers["ChatGPT-Account-Id"] = account_id + headers[CHATGPT_ACCOUNT_ID_HEADER] = account_id return headers +def decode_jwt_claims(token: str) -> JsonObject: + try: + parts: Final = token.split(".") + if len(parts) < 2: + return _EMPTY_CLAIMS + payload_b64: Final = parts[1] + "=" * (-len(parts[1]) % 4) + payload_bytes: Final = base64.urlsafe_b64decode(payload_b64) + return _JSON_OBJECT_ADAPTER.validate_python(json.loads(payload_bytes.decode("utf-8"))) + except Exception: + return _EMPTY_CLAIMS + + +def extract_chatgpt_account_id(token: str | None) -> str | None: + if not token: + return None + auth_claims: Final = decode_jwt_claims(token).get(OPENAI_AUTH_CLAIM_KEY) + if not isinstance(auth_claims, dict): + return None + account_id: Final = auth_claims.get("chatgpt_account_id") + return account_id if isinstance(account_id, str) and account_id else None + + +def _bearer_token(value: str) -> str: + scheme, separator, token = value.strip().partition(" ") + if not separator: + return scheme + return token.strip() if scheme.lower() == "bearer" else "" + + +def is_chatgpt_oauth_key(value: str | None) -> bool: + if value is None: + return False + return OPENAI_AUTH_CLAIM_KEY in decode_jwt_claims(_bearer_token(value)) + + +@dataclass(frozen=True, slots=True) +class ChatGPTClientCredential: + access_token: str + account_id: str | None + + +def _header_value(headers: Mapping[str, str], name: str) -> str | None: + return next((value for key, value in headers.items() if key.lower() == name), None) + + +def get_chatgpt_client_credential(headers: Mapping[str, str]) -> ChatGPTClientCredential | None: + auth_header: Final = _header_value(headers, "authorization") + if auth_header is None or not is_chatgpt_oauth_key(auth_header): + return None + access_token: Final = _bearer_token(auth_header) + account_id: Final = _header_value(headers, CHATGPT_ACCOUNT_ID_HEADER.lower()) or extract_chatgpt_account_id( + access_token + ) + return ChatGPTClientCredential(access_token=access_token, account_id=account_id) + + +def select_chatgpt_client_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + if get_chatgpt_client_credential(headers) is None: + return MappingProxyType({}) + return MappingProxyType( + { + key: value + for key, value in headers.items() + if (key.lower() == "authorization" and is_chatgpt_oauth_key(value)) + or key.lower() == CHATGPT_ACCOUNT_ID_HEADER.lower() + } + ) + + +def strip_chatgpt_client_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + return MappingProxyType( + { + key: value + for key, value in headers.items() + if key.lower() not in ("authorization", CHATGPT_ACCOUNT_ID_HEADER.lower()) + } + ) + + +def get_chatgpt_request_credential( + params: "Mapping[str, object] | GenericLiteLLMParams | None", +) -> ChatGPTClientCredential | None: + normalized: Final = _normalize_litellm_params(params) + scoped: Final = normalized.get("provider_specific_header") + extra: Final = normalized.get("extra_headers") + headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=_SCOPED_HEADERS.validate_python(scoped), custom_llm_provider="chatgpt" + ) + return get_chatgpt_client_credential(headers) or get_chatgpt_client_credential( + _STRING_HEADERS.validate_python(extra) if extra is not None else MappingProxyType({}) + ) + + +def chatgpt_client_api_base(api_base: str | None) -> str: + base: Final = api_base or CHATGPT_API_BASE + parsed: Final = urlsplit(base) + if ( + parsed.scheme != "https" + or parsed.hostname != "chatgpt.com" + or parsed.port not in (None, 443) + or parsed.username + or parsed.password + ): + raise ChatGPTAuthError( + status_code=400, message="Client ChatGPT credentials require the official HTTPS ChatGPT endpoint" + ) + return base + + def get_chatgpt_default_instructions() -> str: return os.getenv("CHATGPT_DEFAULT_INSTRUCTIONS") or CHATGPT_DEFAULT_INSTRUCTIONS -def _normalize_litellm_params(litellm_params: Any | None) -> dict: - if litellm_params is None: - return {} - if isinstance(litellm_params, dict): - return litellm_params - if hasattr(litellm_params, "model_dump"): - try: - return litellm_params.model_dump() - except Exception: - return {} - if hasattr(litellm_params, "dict"): - try: - return litellm_params.dict() - except Exception: - return {} - return {} +def _normalize_litellm_params(litellm_params: object) -> Mapping[str, object]: + if isinstance(litellm_params, Mapping): + return _PARAMS.validate_python(litellm_params) + serializer: Final = getattr(litellm_params, "model_dump", None) or getattr(litellm_params, "dict", None) + if not callable(serializer): + return MappingProxyType({}) + try: + return _PARAMS.validate_python(serializer()) + except Exception: + return MappingProxyType({}) def get_chatgpt_session_id(litellm_params: object) -> str | None: @@ -274,8 +401,8 @@ def get_chatgpt_session_id(litellm_params: object) -> str | None: value = params.get(key) if value: return str(value) - metadata: Final = params.get("metadata") - if isinstance(metadata, dict): + metadata: Final = _normalize_litellm_params(params.get("metadata")) + if metadata: value = metadata.get("session_id") if value: return str(value) diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 3e7db2981ef..2352de35b5b 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hid from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( safe_convert_created_field, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.sse_output_recovery import ( @@ -24,13 +25,15 @@ from litellm.types.llms.openai import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders -from ..authenticator import Authenticator +from ..authenticator import Authenticator, get_cached_authenticator, get_chatgpt_auth_file from ..common_utils import ( CHATGPT_API_BASE, GetAccessTokenError, ensure_chatgpt_session_id, + get_chatgpt_client_credential, get_chatgpt_default_headers, get_chatgpt_default_instructions, + strip_chatgpt_client_credential_headers, ) if TYPE_CHECKING: @@ -48,14 +51,35 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.CHATGPT + def _resolve_authenticator(self, litellm_params: GenericLiteLLMParams | None) -> Authenticator: + auth_file: Final = get_chatgpt_auth_file(litellm_params) + return get_cached_authenticator(auth_file) if auth_file else self.authenticator + def validate_environment( self, - headers: dict, + headers: dict[str, str], model: str, litellm_params: GenericLiteLLMParams | None, - ) -> dict: + ) -> dict[str, str]: + session_id: Final = ensure_chatgpt_session_id(litellm_params) + client_credential: Final = get_chatgpt_client_credential(headers) + if client_credential is not None: + return { + **get_chatgpt_default_headers(client_credential.access_token, client_credential.account_id, session_id), + **strip_chatgpt_client_credential_headers(headers), + } + authenticator: Final = self._resolve_authenticator(litellm_params) + safe_headers: Final = strip_chatgpt_client_credential_headers(headers) + configured_api_key: Final = litellm_params.api_key if litellm_params is not None else None + if configured_api_key and not get_chatgpt_auth_file(litellm_params): + return { + **get_chatgpt_default_headers( + configured_api_key, authenticator.get_account_id(configured_api_key), session_id + ), + **safe_headers, + } try: - access_token: Final = self.authenticator.get_access_token() + access_token: Final = authenticator.get_access_token() except GetAccessTokenError as e: raise AuthenticationError( model=model, @@ -63,10 +87,9 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): message=str(e), ) - account_id: Final = self.authenticator.get_account_id() - session_id: Final = ensure_chatgpt_session_id(litellm_params) + account_id: Final = authenticator.get_account_id() default_headers: Final = get_chatgpt_default_headers(access_token, account_id, session_id) - return {**default_headers, **headers} + return {**default_headers, **safe_headers} def transform_responses_api_request( self, @@ -78,7 +101,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> dict: request: Final = super().transform_responses_api_request( model, - input, + [{"role": "user", "content": input}] if isinstance(input, str) else input, response_api_optional_request_params, litellm_params, headers, @@ -253,8 +276,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: dict, ) -> str: api_base = api_base or self.authenticator.get_api_base() or CHATGPT_API_BASE - api_base = api_base.rstrip("/") - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def supports_native_websocket(self) -> bool: """ChatGPT does not support native WebSocket for Responses API""" diff --git a/litellm/llms/clarifai/chat/transformation.py b/litellm/llms/clarifai/chat/transformation.py index 8c2fd9e0aff..d07ae5e4215 100644 --- a/litellm/llms/clarifai/chat/transformation.py +++ b/litellm/llms/clarifai/chat/transformation.py @@ -23,6 +23,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class ClarifaiConfig(OpenAIGPTConfig): @@ -114,7 +115,7 @@ class ClarifaiConfig(OpenAIGPTConfig): ) ## Reponse try: - completion_response: Final = raw_response.json() + completion_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise OpenAIError( status_code=raw_response.status_code, diff --git a/litellm/llms/claude_code/harness/transformation.py b/litellm/llms/claude_code/harness/transformation.py index f75d86c7e63..8252d2dd461 100644 --- a/litellm/llms/claude_code/harness/transformation.py +++ b/litellm/llms/claude_code/harness/transformation.py @@ -30,6 +30,7 @@ from litellm.harness.types import ( ToolCall, ToolResult, ) +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.base_llm.harness.transformation import ( BaseCLIHarnessConfig, HarnessSessionSetup, @@ -231,12 +232,12 @@ def _system_events(event: Mapping[str, object], state: ClaudeCodeStreamState) -> return event_list() -def _record_result(event: Mapping[str, Any], state: ClaudeCodeStreamState) -> Sequence[Event]: +def _record_result(event: Mapping[str, object], state: ClaudeCodeStreamState) -> Sequence[Event]: state.result_seen = True state.is_error = bool(event.get("is_error", False)) result = event.get("result") state.result_text = result if isinstance(result, str) else None - state.errors = [str(e) for e in event.get("errors") or ()] # mutable-ok: mirrors the JSON errors array + state.errors = [str(e) for e in iter_items(event.get("errors") or ())] # mutable-ok: mirrors the JSON errors array state.structured_output = event.get("structured_output") if event.get("session_id"): state.session_id = str(event["session_id"]) diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index ef11e51a5eb..2bbfd2010cc 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -3,6 +3,7 @@ from typing import Final import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.secret_managers.main import ( @@ -57,9 +58,9 @@ class CloudflareChatConfig(OpenAIGPTConfig): trimmed: Final = api_base.rstrip("/") if trimmed.endswith("/ai/run"): verbose_logger.warning( - "Cloudflare api_base ending in '/ai/run' is the legacy Workers AI path and no longer serves OpenAI-compatible requests; rewriting to the '/ai/v1' endpoint" + "Cloudflare chat completions require the OpenAI-compatible '/ai/v1' endpoint; rewriting api_base ending in '/ai/run' to '/ai/v1'. Decision models continue to use the native Workers AI '/ai/run/{model}' endpoint." ) - return f"{trimmed[: -len('/ai/run')]}/ai/v1" + return join_api_base(trimmed[: -len("/ai/run")], "ai/v1") return api_base def validate_environment( diff --git a/litellm/llms/cloudflare/common_utils.py b/litellm/llms/cloudflare/common_utils.py index 9f1293cb9c4..02855a879a2 100644 --- a/litellm/llms/cloudflare/common_utils.py +++ b/litellm/llms/cloudflare/common_utils.py @@ -3,6 +3,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import ( get_secret_str, normalize_nonempty_secret_str, @@ -34,10 +35,10 @@ def workers_ai_request_model(model: str) -> str: def workers_ai_run_url(api_base: str, model: str) -> str: normalized_api_base: Final = api_base.rstrip("/") if normalized_api_base.endswith("/ai/v1"): - return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}" + return join_api_base(normalized_api_base.removesuffix("/ai/v1"), f"ai/run/{model}") if normalized_api_base.endswith("/ai/run"): return f"{normalized_api_base}/{model}" - return f"{normalized_api_base}/ai/run/{model}" + return join_api_base(normalized_api_base, f"ai/run/{model}") def unwrap_workers_ai_result(payload: object) -> object: diff --git a/litellm/llms/cloudflare/decisions/transformation.py b/litellm/llms/cloudflare/decisions/transformation.py index 7c776367886..7575e8bce21 100644 --- a/litellm/llms/cloudflare/decisions/transformation.py +++ b/litellm/llms/cloudflare/decisions/transformation.py @@ -20,6 +20,8 @@ class CloudflareDecisionsConfig(BaseDecisionsConfig): inlines_remote_images = True def get_default_api_base(self) -> str | None: + # Decision models (Clef) require native /ai/run/{model} requests with + # state and questions; /ai/v1/chat/completions only accepts chat requests. return workers_ai_default_api_base() def missing_api_base_message(self, custom_llm_provider: str) -> str: diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 08799c3a0a5..6bebfe2d444 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -2,7 +2,7 @@ ## handler file for TextCompletionCodestral Integration - https://codestral.com/ import json -from collections.abc import Callable +from collections.abc import Callable, Mapping from functools import partial from typing import Final, Protocol @@ -268,7 +268,7 @@ class CodestralTextCompletion: prompt = prompt_factory(model=model, messages=messages) ## Load Config - config: Final = litellm.CodestralTextCompletionConfig.get_config() + config: Final[Mapping[str, object]] = litellm.CodestralTextCompletionConfig.get_config() for k, v in config.items(): if ( k not in optional_params diff --git a/litellm/llms/codex/harness/transformation.py b/litellm/llms/codex/harness/transformation.py index 4512e39189e..d7388bd23fc 100644 --- a/litellm/llms/codex/harness/transformation.py +++ b/litellm/llms/codex/harness/transformation.py @@ -15,7 +15,7 @@ import re from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass, field from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from pydantic import ConfigDict, TypeAdapter @@ -31,6 +31,7 @@ from litellm.harness.types import ( ToolCall, ToolResult, ) +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.base_llm.harness.transformation import ( BaseCLIHarnessConfig, HarnessSessionSetup, @@ -77,13 +78,13 @@ class CodexStreamState: started: set[str] = field(default_factory=set) # mutable-ok: parser records announced tool items -def _tool_input(item: Mapping[str, Any]) -> tuple[str, str, Mapping[str, object], bool]: +def _tool_input(item: Mapping[str, object]) -> tuple[str, str, Mapping[str, object], bool]: """(normalized name, native name, input, builtin) for a tool-like item.""" item_type = item.get("type") if item_type == "command_execution": return "bash", "command_execution", MappingProxyType({"command": item.get("command", "")}), True if item_type == "file_change": - changes: Final = list(item.get("changes") or ()) # mutable-ok: JSON array, as codex reports it + changes: Final = list(iter_items(item.get("changes") or ())) # mutable-ok: JSON array, as codex reports it return "edit", "apply_patch", MappingProxyType({"changes": changes}), True if item_type == "web_search": return "web_search", "web_search", MappingProxyType({"query": item.get("query", "")}), True diff --git a/litellm/llms/cohere/chat/transformation.py b/litellm/llms/cohere/chat/transformation.py index fdf44d5333e..e7fadd7f375 100644 --- a/litellm/llms/cohere/chat/transformation.py +++ b/litellm/llms/cohere/chat/transformation.py @@ -1,6 +1,6 @@ import json import time -from collections.abc import AsyncIterator, Iterator, Sequence +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx @@ -221,7 +221,8 @@ class CohereChatConfig(BaseConfig): headers: dict, ) -> dict: ## Load Config - for k, v in litellm.CohereChatConfig.get_config().items(): + config: Final[Mapping[str, object]] = litellm.CohereChatConfig.get_config() + for k, v in config.items(): if ( k not in optional_params ): # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index b6d94d4853a..7d9c8de7a07 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -6,6 +6,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -14,6 +15,7 @@ from litellm.types.rerank import OptionalRerankParams, RerankRequest, RerankResp from ..common_utils import CohereError _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class CohereRerankConfig(BaseRerankConfig): @@ -31,11 +33,7 @@ class CohereRerankConfig(BaseRerankConfig): optional_params: dict | None = None, ) -> str: if api_base: - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/rerank"): - api_base = f"{api_base}/v1/rerank" - return api_base + return ensure_api_route(api_base, "v1/rerank") return "https://api.cohere.ai/v1/rerank" def get_supported_cohere_rerank_params(self, model: str) -> list: @@ -147,7 +145,7 @@ class CohereRerankConfig(BaseRerankConfig): No transformation required, litellm follows cohere API response format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raise CohereError(message=raw_response.text, status_code=raw_response.status_code) diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index e79357bf52f..83596f90d1f 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -1,5 +1,6 @@ -from typing import Any, Final +from typing import Final +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.types.rerank import OptionalRerankParams, RerankRequest @@ -19,11 +20,7 @@ class CohereRerankV2Config(CohereRerankConfig): optional_params: dict | None = None, ) -> str: if api_base: - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/v2/rerank"): - api_base = f"{api_base}/v2/rerank" - return api_base + return ensure_api_route(api_base, "v2/rerank") return "https://api.cohere.ai/v2/rerank" def get_supported_cohere_rerank_params(self, model: str) -> list: @@ -42,7 +39,7 @@ class CohereRerankV2Config(CohereRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/llms/cometapi/chat/transformation.py b/litellm/llms/cometapi/chat/transformation.py index 26050700efc..ae04889bc97 100644 --- a/litellm/llms/cometapi/chat/transformation.py +++ b/litellm/llms/cometapi/chat/transformation.py @@ -6,10 +6,11 @@ Documentation: [CometAPI Documentation Link] """ from collections.abc import AsyncIterator, Iterator -from typing import Any, Final +from typing import Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -112,16 +113,14 @@ class CometAPIConfig(OpenAIGPTConfig): # Ensure we include /v1 prefix when missing if api_base.endswith("/v1"): - return f"{api_base}/{endpoint}" - if api_base.endswith("/v1/"): - return f"{api_base}{endpoint}" + return join_api_base(api_base, endpoint) # If user provided https://api.cometapi.com, add /v1 if api_base == "https://api.cometapi.com": - return f"{api_base}/v1/{endpoint}" + return join_api_base(api_base, f"v1/{endpoint}") # Generic fallback: if '/v1' not in path, add it if "/v1" not in api_base.split("//", 1)[-1]: - return f"{api_base}/v1/{endpoint}" - return f"{api_base}/{endpoint}" + return join_api_base(api_base, f"v1/{endpoint}") + return join_api_base(api_base, endpoint) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: """ @@ -138,7 +137,7 @@ class CometAPIConfig(OpenAIGPTConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "CometAPIChatCompletionStreamingHandler": """ Get model response iterator for streaming responses """ diff --git a/litellm/llms/cometapi/embed/transformation.py b/litellm/llms/cometapi/embed/transformation.py index 60f781be261..e3b5c722664 100644 --- a/litellm/llms/cometapi/embed/transformation.py +++ b/litellm/llms/cometapi/embed/transformation.py @@ -9,6 +9,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -44,9 +45,7 @@ class CometAPIEmbeddingConfig(BaseEmbeddingConfig): """ Get the complete URL for the CometAPI embedding endpoint. """ - api_base = "https://api.cometapi.com/v1" if api_base is None else api_base.rstrip("/") - complete_url: Final = f"{api_base}/embeddings" - return complete_url + return join_api_base("https://api.cometapi.com/v1" if api_base is None else api_base, "embeddings") def validate_environment( self, diff --git a/litellm/llms/cometapi/image_generation/transformation.py b/litellm/llms/cometapi/image_generation/transformation.py index 567dbfd308e..ffaabbb28cd 100644 --- a/litellm/llms/cometapi/image_generation/transformation.py +++ b/litellm/llms/cometapi/image_generation/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -77,16 +78,13 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete url for the request """ - complete_url: str = ( + base_url: Final = ( api_base or get_secret_str("COMETAPI_BASE_URL") or get_secret_str("COMETAPI_API_BASE") or self.DEFAULT_BASE_URL ) - - complete_url = complete_url.rstrip("/") - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" - return complete_url + return join_api_base(base_url, self.IMAGE_GENERATION_ENDPOINT) def validate_environment( self, diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index a02db2338d8..87fc61985e6 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -3,10 +3,10 @@ CompactifAI chat completion transformation """ from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import httpx -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import OpenAIError @@ -25,11 +25,32 @@ else: LiteLLMLoggingObj = Any +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + +class _CompactifAIToolCall(TypedDict): + function: ReadOnly[Mapping[str, str]] + + +class _CompactifAIMessage(TypedDict, total=False): + content: str | None # writable-ok: json mode moves a lone tool call's arguments into the content in place + tool_calls: Sequence[_CompactifAIToolCall] | None # writable-ok: json mode clears the tool call it moved + + +class _CompactifAIChoice(TypedDict, total=False): + message: ReadOnly[_CompactifAIMessage] + + class CompactifAIResponseFields(TypedDict, total=False): """The chat completion fields of a CompactifAI response body.""" id: ReadOnly[str] - choices: ReadOnly[Sequence[Mapping[str, object]]] + choices: ReadOnly[Required[Sequence[_CompactifAIChoice]]] created: ReadOnly[int] model: ReadOnly[str] system_fingerprint: ReadOnly[str | None] @@ -82,22 +103,19 @@ class CompactifAIChatConfig(OpenAIGPTConfig): ) ## RESPONSE OBJECT - response_json: Final = raw_response.json() + response_json: Final = cast( # cast-ok: the body passes through unvalidated, as before it was typed + CompactifAIResponseFields, _json_body(raw_response) + ) # Handle JSON mode if needed if json_mode: for choice in response_json["choices"]: message = choice.get("message") - if message and message.get("tool_calls"): - # Convert tool calls to content for JSON mode - tool_calls = message.get("tool_calls", []) - if len(tool_calls) == 1: - message["content"] = tool_calls[0]["function"].get("arguments", "") - message["tool_calls"] = None + if message and (tool_calls := message.get("tool_calls")) and len(tool_calls) == 1: + message["content"] = tool_calls[0]["function"].get("arguments", "") + message["tool_calls"] = None - response_fields: Final[CompactifAIResponseFields] = response_json - - returned_response: Final = ModelResponse(**response_fields) + returned_response: Final = ModelResponse(**response_json) # Set model name with provider prefix returned_response.model = f"compactifai/{model}" diff --git a/litellm/llms/custom_httpx/async_client_cleanup.py b/litellm/llms/custom_httpx/async_client_cleanup.py index 6ef61b6f58f..3c0cc90f30f 100644 --- a/litellm/llms/custom_httpx/async_client_cleanup.py +++ b/litellm/llms/custom_httpx/async_client_cleanup.py @@ -20,7 +20,7 @@ async def close_litellm_async_clients(): cache_dict: Final = getattr(litellm.in_memory_llm_clients_cache, "cache_dict", {}) - for key, handler in cache_dict.items(): + for handler in cache_dict.values(): # Handle BaseLLMAIOHTTPHandler instances (aiohttp_openai provider) if isinstance(handler, BaseLLMAIOHTTPHandler) and hasattr(handler, "close"): try: diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index cffeb44f12c..0b64885df45 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -20,6 +20,7 @@ from aiohttp import ClientSession, DummyCookieJar, TCPConnector from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport from httpx._types import CertTypes, RequestFiles from httpx._utils import get_environment_proxies +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -566,6 +567,34 @@ async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> N raise MaskedHTTPStatusError(e, message=_text, text=_text) from None +_TIMEOUT_PHASES: Final[Mapping[type[httpx.TimeoutException], tuple[str, str]]] = MappingProxyType( + { + httpx.ConnectTimeout: ("Connection", "connect"), + httpx.ReadTimeout: ("Read", "read"), + httpx.WriteTimeout: ("Write", "write"), + httpx.PoolTimeout: ("Connection pool", "pool"), + } +) +_UNKNOWN_TIMEOUT_PHASE: Final = ("Request", "read") +_TIMEOUT_EXTENSION: Final = TypeAdapter(dict[str, float | None]) + + +def _timeout_phase_seconds(error: httpx.TimeoutException, phase: str) -> float | None: + """The limit httpx enforced for ``phase`` on the request that raised ``error``, from the request itself.""" + return _TIMEOUT_EXTENSION.validate_python(error.request.extensions["timeout"]).get(phase) + + +def _raise_timeout(error: httpx.TimeoutException, started_at: float) -> NoReturn: + label, phase = _TIMEOUT_PHASES.get(type(error), _UNKNOWN_TIMEOUT_PHASE) + seconds: Final = _timeout_phase_seconds(error, phase) + time_taken: Final = round(time.time() - started_at, 3) + raise litellm.Timeout( + message=f"{label} timed out. Timeout passed={seconds}, time taken={time_taken} seconds", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ) + + class HTTPResponseLimitError(ValueError): pass @@ -870,20 +899,7 @@ class AsyncHTTPHandler: finally: await new_client.aclose() except httpx.TimeoutException as e: - end_time: Final = time.time() - time_delta: Final = round(end_time - start_time, 3) - headers = {} - error_response: Final[httpx.Response | None] = getattr(e, "response", None) - if error_response is not None: - for key, value in error_response.headers.items(): - headers[f"response_headers-{key}"] = value - - raise litellm.Timeout( - message=f"Connection timed out. Timeout passed={timeout}, time taken={time_delta} seconds", - model="default-model-name", - llm_provider="litellm-httpx-handler", - headers=headers, - ) + _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: await _raise_masked_async_error(e, stream) except Exception as e: @@ -902,6 +918,7 @@ class AsyncHTTPHandler: follow_redirects: bool | None = None, ): _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT + start_time: Final = time.time() try: if timeout is None: timeout = self.timeout @@ -946,18 +963,7 @@ class AsyncHTTPHandler: finally: await new_client.aclose() except httpx.TimeoutException as e: - headers = {} - error_response: Final[httpx.Response | None] = getattr(e, "response", None) - if error_response is not None: - for key, value in error_response.headers.items(): - headers[f"response_headers-{key}"] = value - - raise litellm.Timeout( - message=f"Connection timed out after {timeout} seconds.", - model="default-model-name", - llm_provider="litellm-httpx-handler", - headers=headers, - ) + _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: await _raise_masked_async_error(e, stream) except Exception as e: @@ -974,6 +980,7 @@ class AsyncHTTPHandler: stream: bool = False, content: _RequestContent | None = None, ): + start_time: Final = time.time() try: if timeout is None: timeout = self.timeout @@ -1012,18 +1019,7 @@ class AsyncHTTPHandler: finally: await new_client.aclose() except httpx.TimeoutException as e: - headers = {} - error_response: Final[httpx.Response | None] = getattr(e, "response", None) - if error_response is not None: - for key, value in error_response.headers.items(): - headers[f"response_headers-{key}"] = value - - raise litellm.Timeout( - message=f"Connection timed out after {timeout} seconds.", - model="default-model-name", - llm_provider="litellm-httpx-handler", - headers=headers, - ) + _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: await _raise_masked_async_error(e, stream) except Exception as e: @@ -1514,6 +1510,7 @@ class HTTPHandler: content: _RequestContent | None = None, logging_obj: LiteLLMLoggingObject | None = None, ): + start_time: Final = time.time() try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) request_data, request_content = _prepare_request_data_and_content(data, content) @@ -1546,12 +1543,8 @@ class HTTPHandler: _anchor_handler_to(response, self) response.raise_for_status() return response - except httpx.TimeoutException: - raise litellm.Timeout( - message=f"Connection timed out after {timeout} seconds.", - model="default-model-name", - llm_provider="litellm-httpx-handler", - ) + except httpx.TimeoutException as e: + _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: _raise_masked_sync_error(e, stream) except Exception as e: @@ -1568,6 +1561,7 @@ class HTTPHandler: timeout: float | httpx.Timeout | None = None, content: _RequestContent | None = None, ): + start_time: Final = time.time() try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) request_data, request_content = _prepare_request_data_and_content(data, content) @@ -1598,12 +1592,8 @@ class HTTPHandler: _anchor_handler_to(response, self) response.raise_for_status() return response - except httpx.TimeoutException: - raise litellm.Timeout( - message=f"Connection timed out after {timeout} seconds.", - model="default-model-name", - llm_provider="litellm-httpx-handler", - ) + except httpx.TimeoutException as e: + _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: _raise_masked_sync_error(e, stream) except Exception as e: @@ -1620,6 +1610,7 @@ class HTTPHandler: timeout: float | httpx.Timeout | None = None, content: _RequestContent | None = None, ): + start_time: Final = time.time() try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) request_data, request_content = _prepare_request_data_and_content(data, content) @@ -1649,12 +1640,8 @@ class HTTPHandler: if stream: _anchor_handler_to(response, self) return response - except httpx.TimeoutException: - raise litellm.Timeout( - message=f"Connection timed out after {timeout} seconds.", - model="default-model-name", - llm_provider="litellm-httpx-handler", - ) + except httpx.TimeoutException as e: + _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: _raise_masked_sync_error(e, stream) except Exception as e: @@ -1671,6 +1658,7 @@ class HTTPHandler: stream: bool = False, content: _RequestContent | None = None, ): + start_time: Final = time.time() try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) request_data, request_content = _prepare_request_data_and_content(data, content) @@ -1701,12 +1689,8 @@ class HTTPHandler: _anchor_handler_to(response, self) response.raise_for_status() return response - except httpx.TimeoutException: - raise litellm.Timeout( - message=f"Connection timed out after {timeout} seconds.", - model="default-model-name", - llm_provider="litellm-httpx-handler", - ) + except httpx.TimeoutException as e: + _raise_timeout(e, start_time) except httpx.HTTPStatusError as e: _raise_masked_sync_error(e, stream) except Exception as e: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ac2ac1fe31b..7b808cf2be2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -49,17 +49,20 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SUBTITLE_RESPONSE_FORMATS, synthesize_subtitle_document, ) -from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, 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.get_llm_provider_logic import get_provider_affinity_header from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.litellm_core_utils.realtime_errors import ( close_after_upstream_handshake_refusal, realtime_error_event, websocket_close_reason, ) from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base +from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -195,6 +198,26 @@ from litellm.utils import ( ) +def _merge_headers_case_insensitive( + headers: Mapping[str, object], extra_headers: Mapping[str, object] | None +) -> dict[str, str]: # mutable-ok: legacy HTTP handler methods require a dict of headers + overridden: Final = frozenset(name.lower() for name in extra_headers or {}) + return { + **{name: str(value) for name, value in headers.items() if name.lower() not in overridden}, + **{name: str(value) for name, value in (extra_headers or {}).items()}, + } + + +def _client_request_headers( + response_api_optional_request_params: Mapping[str, object], + extra_headers: Mapping[str, object] | None, +) -> dict[str, str]: # mutable-ok: provider validate_environment methods require a dict + from litellm.litellm_core_utils.core_helpers import as_str_mapping + + optional_headers: Final = as_str_mapping(response_api_optional_request_params.get("extra_headers")) + return _merge_headers_case_insensitive(optional_headers or {}, extra_headers) + + def _rust_responses_websocket_enabled( custom_llm_provider: str | None, ) -> bool: @@ -409,6 +432,27 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False +def _run_responses_deployment_hook( + *, + model: str, + request_input: str | ResponseInputParam, + custom_llm_provider: str, + optional_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, +) -> Mapping[str, object]: + return run_async_function( + async_pre_call_deployment_hook, + { + **dict(litellm_params), + **optional_params, + "model": model, + "input": request_input, + "custom_llm_provider": custom_llm_provider, + }, + CallTypes.responses.value, + ) + + def _mask_presigned_request_headers(transformed_request: bytes | str | dict) -> bytes | str | dict: """A pre-signed request carries its auth inside its own ``headers`` key, which logging treats as request body (only the top-level headers channel gets masked), @@ -474,6 +518,7 @@ _DECODED_BODY_STALE_HEADERS: Final[frozenset[str]] = frozenset({"content-encodin _OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object]) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _JSON_OBJECTS: Final = TypeAdapter(list[dict[str, object]]) +_DECODED_JSON: Final = TypeAdapter(object) def _decoded_body_headers(response: httpx.Response) -> httpx.Headers: @@ -974,7 +1019,7 @@ class BaseLLMHTTPHandler: original_data: dict, model: str, messages: list, - logging_obj, + logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, timeout: float | httpx.Timeout, @@ -1028,6 +1073,7 @@ class BaseLLMHTTPHandler: sync_stream=True, json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), + logging_obj=logging_obj, ) if isinstance(completion_stream, BaseModelResponseIterator): completion_stream.http_response = response @@ -1170,6 +1216,7 @@ class BaseLLMHTTPHandler: sync_stream=False, json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), + logging_obj=logging_obj, ) if isinstance(completion_stream, BaseModelResponseIterator): completion_stream.http_response = response @@ -2435,13 +2482,13 @@ class BaseLLMHTTPHandler: # Also check for extra_headers in kwargs (from config or direct calls) extra_headers_from_kwargs: Final = kwargs.get("extra_headers", None) # Merge all header sources: forwarded < extra_headers < provider_specific - merged_headers: Final = {} - if forwarded_headers: - merged_headers.update(forwarded_headers) - if extra_headers_from_kwargs: - merged_headers.update(extra_headers_from_kwargs) - if provider_specific_headers: - merged_headers.update(provider_specific_headers) + merged_headers: Final[Mapping[str, object]] = MappingProxyType( + { + **(forwarded_headers or {}), + **(extra_headers_from_kwargs or {}), + **(provider_specific_headers or {}), + } + ) request_params: Final = reduce( delete_nested_value, litellm_params.get("additional_drop_params") or (), @@ -2452,11 +2499,19 @@ class BaseLLMHTTPHandler: and anthropic_messages_provider_config.should_filter_anthropic_beta_headers() ): await aload_beta_headers_config() + affinity_api_base: Final = ( + AnthropicModelInfo.get_api_base(api_base) if custom_llm_provider == "anthropic" else api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) ( headers, api_base, ) = await anthropic_messages_provider_config.avalidate_anthropic_messages_environment( - headers=merged_headers or {}, + headers=( + add_provider_affinity_header(merged_headers, litellm_params, default_affinity_header) + if default_affinity_header is not None + else dict(merged_headers) + ), model=model, messages=messages, optional_params=request_params, @@ -2754,16 +2809,12 @@ class BaseLLMHTTPHandler: litellm_params, ) - modified_kwargs: Final = run_async_function( - async_pre_call_deployment_hook, - { - **dict(litellm_params), - **response_api_optional_request_params, - "model": model, - "input": input, - "custom_llm_provider": custom_llm_provider, - }, - CallTypes.responses.value, + modified_kwargs: Final = _run_responses_deployment_hook( + model=model, + request_input=input, + custom_llm_provider=custom_llm_provider, + optional_params=response_api_optional_request_params, + litellm_params=litellm_params, ) if modified_kwargs is None: return ( @@ -2875,14 +2926,18 @@ class BaseLLMHTTPHandler: else: sync_httpx_client = client - headers = responses_api_provider_config.validate_environment( - headers=response_api_optional_request_params.get("extra_headers", {}) or {}, - model=model, - litellm_params=litellm_params, - ) - - if extra_headers: - headers.update(extra_headers) + headers = { # rebind-ok: validation replaces the method's incoming headers + **responses_api_provider_config.validate_environment( + headers=( + _client_request_headers(response_api_optional_request_params, extra_headers) + if custom_llm_provider == "chatgpt" + else response_api_optional_request_params.get("extra_headers", {}) or {} + ), + model=model, + litellm_params=litellm_params, + ), + **((extra_headers or {}) if custom_llm_provider != "chatgpt" else {}), + } # Check if streaming is requested stream = response_api_optional_request_params.get("stream", False) @@ -3073,11 +3128,15 @@ class BaseLLMHTTPHandler: validated_headers: Final = { **await asyncio.to_thread( responses_api_provider_config.validate_environment, - headers=response_api_optional_request_params.get("extra_headers", {}) or {}, + headers=( + _client_request_headers(response_api_optional_request_params, extra_headers) + if custom_llm_provider == "chatgpt" + else response_api_optional_request_params.get("extra_headers", {}) or {} + ), model=model, litellm_params=litellm_params, ), - **(extra_headers or {}), + **((extra_headers or {}) if custom_llm_provider != "chatgpt" else {}), } # Check if streaming is requested @@ -3742,7 +3801,7 @@ class BaseLLMHTTPHandler: response: httpx.Response, upload_url_location: str, upload_url_key: str = "upload_url", - ) -> tuple[str | None, dict | None]: + ) -> tuple[object, Mapping[str, object] | None]: """ Extract upload URL from initial file creation response. @@ -3762,8 +3821,10 @@ class BaseLLMHTTPHandler: else: # Response body style (e.g., Manus, S3 presigned URLs) try: - response_data: Final = response.json() - upload_url = response_data.get(upload_url_key) + response_data: Final = as_str_mapping(_DECODED_JSON.validate_python(response.json())) + if response_data is None: + return None, None + upload_url: Final = response_data.get(upload_url_key) return upload_url, response_data if upload_url else None except Exception: return None, None @@ -3929,7 +3990,7 @@ class BaseLLMHTTPHandler: elif isinstance(transformed_request, dict) and "file" in transformed_request: # Handle multipart form-data uploads (e.g., Anthropic Files API) # The dict contains tuples suitable for httpx's `files` parameter - file_request: Final = cast(dict[str, Any], transformed_request) + file_request: Final = cast(dict[str, object], transformed_request) upload_response = sync_httpx_client.post( url=api_base, headers=headers, @@ -6937,7 +6998,7 @@ class BaseLLMHTTPHandler: headers={}, model=model or "", api_key=api_key ) else: - url = f"{api_base.rstrip('/')}/v1/realtime/{endpoint}" + url = join_api_base(api_base, f"v1/realtime/{endpoint}") headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", @@ -7006,7 +7067,7 @@ class BaseLLMHTTPHandler: url = provider_config.get_realtime_calls_url(api_base=api_base, model=model or "", api_version=api_version) headers: dict[str, object] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) else: - url = f"{api_base.rstrip('/')}/v1/realtime/calls" + url = join_api_base(api_base, "v1/realtime/calls") headers = { "Authorization": f"Bearer {openai_ephemeral_key}", } @@ -10858,7 +10919,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) logging_obj.pre_call( input="", @@ -10923,7 +10984,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) logging_obj.pre_call( input="", @@ -11126,7 +11187,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) request_body: Final[dict[str, object]] = dict(vector_store_update_optional_params) metadata: Final = vector_store_update_optional_params.get("metadata") @@ -11205,7 +11266,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) request_body: Final[dict[str, object]] = dict(vector_store_update_optional_params) metadata: Final = vector_store_update_optional_params.get("metadata") @@ -11273,7 +11334,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) logging_obj.pre_call( input="", @@ -11335,7 +11396,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) logging_obj.pre_call( input="", diff --git a/litellm/llms/custom_httpx/mock_transport.py b/litellm/llms/custom_httpx/mock_transport.py index 0322cceec0f..6acadb8e98a 100644 --- a/litellm/llms/custom_httpx/mock_transport.py +++ b/litellm/llms/custom_httpx/mock_transport.py @@ -12,6 +12,9 @@ import uuid from typing import Final import httpx +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import decoded_json_object # --------------------------------------------------------------------------- # Pre-built response templates @@ -22,7 +25,7 @@ def _mock_id() -> str: return f"chatcmpl-mock-{uuid.uuid4().hex[:8]}" -def _chat_completion_json(model: str) -> dict: +def _chat_completion_json(model: object) -> dict: """Return a minimal valid ChatCompletion object.""" return { "id": _mock_id(), @@ -55,6 +58,8 @@ _JSON_HEADERS: Final = { "content-type": "application/json", } +_DECODED_JSON: Final = TypeAdapter(object) + class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport): """ @@ -64,13 +69,13 @@ class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport): """ @staticmethod - def _parse_request(request: httpx.Request) -> tuple[str, bool]: + def _parse_request(request: httpx.Request) -> tuple[object, bool]: """Extract model from the request body.""" try: - body: Final = json.loads(request.content) + body: Final = _DECODED_JSON.validate_python(json.loads(request.content)) except (json.JSONDecodeError, ValueError): return ("mock-model", False) - model: Final = body.get("model", "mock-model") + model: Final = decoded_json_object(body).get("model", "mock-model") return (model, False) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 9a1b381a2ba..966d1a01b8d 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -5,6 +5,7 @@ Translates from OpenAI's `/v1/chat/completions` to DashScope's `/v1/chat/complet from collections.abc import Coroutine from typing import Final, Literal, overload +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -82,7 +83,4 @@ class DashScopeChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default DashScope /chat/completions endpoint. """ - resolved_api_base: Final = self._resolve_chat_api_base(api_base).rstrip("/") - if resolved_api_base.endswith("/chat/completions"): - return resolved_api_base - return f"{resolved_api_base}/chat/completions" + return ensure_api_route(self._resolve_chat_api_base(api_base), "chat/completions") diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 98504f54533..55d6fb1a0e2 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -7,6 +7,7 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate. See https://help.aliyun.com/zh/model-studio/billing-for-model-studio """ +from collections.abc import Mapping from dataclasses import dataclass, replace from datetime import datetime from typing import Final @@ -54,10 +55,9 @@ def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: ) -def _flat_rate(model_info: ModelInfo, cost_key: str, fallback_cost_key: str) -> float: - value: Final = model_info.get(cost_key) +def _flat_rate(value: float | None, fallback_value: float | None) -> float: if value is None: - return float(model_info.get(fallback_cost_key) or 0.0) + return float(fallback_value or 0.0) return float(value) @@ -65,14 +65,18 @@ def _flat_rates(model_info: ModelInfo) -> TokenRates: reasoning_rate: Final = model_info.get("output_cost_per_reasoning_token") return TokenRates( input_rate=float(model_info.get("input_cost_per_token") or 0.0), - cache_read_rate=_flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token"), - cache_creation_rate=_flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token"), + cache_read_rate=_flat_rate( + model_info.get("cache_read_input_token_cost"), model_info.get("input_cost_per_token") + ), + cache_creation_rate=_flat_rate( + model_info.get("cache_creation_input_token_cost"), model_info.get("input_cost_per_token") + ), output_rate=float(model_info.get("output_cost_per_token") or 0.0), reasoning_rate=None if reasoning_rate is None else float(reasoning_rate), ) -def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates: +def _tier_rates(model_info: ModelInfo, tier: Mapping[str, object]) -> TokenRates: # A tier that declares output rates keeps the request on them, all-or-nothing. A tier table # spelling out only input rates would serve every completion for free, so there the model's # own output rates stand in diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index e5a987ca1ad..d0c2b4dff0c 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -16,6 +16,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -98,10 +99,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - base: Final = self._resolve_embedding_api_base(api_base).rstrip("/") - if base.endswith("/embeddings"): - return base - return f"{base}/embeddings" + return ensure_api_route(self._resolve_embedding_api_base(api_base), "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index 6a172e1e074..d1faa6c6cd0 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -49,6 +49,7 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) DEFAULT_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation" @@ -187,7 +188,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): ) try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Failed to parse DashScope image generation response: {e}", diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 4ec7d697f02..48ce4b5a07b 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -34,6 +34,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -94,10 +95,10 @@ class DashScopeRerankConfig(BaseRerankConfig): return cleaned if cleaned.endswith("/v1"): - return f"{cleaned}/reranks" + return join_api_base(cleaned, "reranks") # Unknown base: append /reranks rather than silently ignoring the caller's api_base. - return f"{cleaned}/reranks" + return join_api_base(cleaned, "reranks") def validate_environment( self, diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 1db91cbb127..5c738423a5a 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -16,7 +16,7 @@ from pydantic import BaseModel, Field, StrictBool, StrictInt, StrictStr, TypeAda import litellm from litellm.exceptions import UnsupportedParamsError -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.videos.async_task_utils import ( input_reference_url, @@ -708,7 +708,7 @@ class DashScopeVideoConfig(BaseVideoConfig): if parameters: request_data["parameters"] = parameters - return request_data, (), f"{api_base}{DASHSCOPE_VIDEO_SYNTHESIS_PATH}" + return request_data, (), join_api_base(api_base, DASHSCOPE_VIDEO_SYNTHESIS_PATH) @staticmethod def _reference_input( @@ -773,7 +773,7 @@ class DashScopeVideoConfig(BaseVideoConfig): def _task_url(video_id: str, api_base: str) -> str: original_task_id: Final = extract_original_video_id(video_id) encoded_task_id: Final = encode_url_path_segment(original_task_id, field_name="video_id") - return f"{api_base}{DASHSCOPE_TASKS_PATH}/{encoded_task_id}" + return join_api_base(api_base, f"{DASHSCOPE_TASKS_PATH}/{encoded_task_id}") def transform_video_status_retrieve_response( self, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index d46d04578cc..48118c23ed7 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_name_from_message, strip_prefill_prefix_field, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import ( type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # base_utils helper returns an untyped dict @@ -339,10 +340,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): stream: bool | None = None, ) -> str: use_ai_gateway: Final = self.uses_ai_gateway(model) - api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway) - url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base - complete_url: Final = f"{url_base}/chat/completions" - return complete_url + return join_api_base(self._get_api_base(api_base, use_ai_gateway=use_ai_gateway), "chat/completions") def get_supported_openai_params(self, model: str | None = None) -> list: return [ @@ -824,6 +822,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): sync_stream: bool, json_mode: bool | None, optional_params: Mapping[str, object], + logging_obj: LiteLLMLoggingObj | None = None, ) -> "DatabricksChatResponseIterator": return self.get_model_response_iterator( streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 3112b2458b4..cd9432ac36f 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -20,6 +20,7 @@ from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -278,10 +279,10 @@ class DatabricksBase: if normalized_api_base.endswith("/ai-gateway/mlflow/v1"): return normalized_api_base if normalized_api_base.endswith("/serving-endpoints"): - return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1" + return join_api_base(normalized_api_base.removesuffix("/serving-endpoints"), "ai-gateway/mlflow/v1") api_base_parts: Final = urlsplit(normalized_api_base) if api_base_parts.path in ("", "/"): - return f"{normalized_api_base}/ai-gateway/mlflow/v1" + return join_api_base(normalized_api_base, "ai-gateway/mlflow/v1") return api_base def get_oauth_m2m_token( @@ -319,7 +320,7 @@ class DatabricksBase: api_base_parts: Final = urlsplit(api_base) workspace_url: Final = urlunsplit((api_base_parts.scheme, api_base_parts.netloc, "", "", "")) - token_url: Final = f"{workspace_url}/oidc/v1/token" + token_url: Final = join_api_base(workspace_url, "oidc/v1/token") try: response: Final = requests.post( @@ -477,5 +478,6 @@ class DatabricksBase: # Debug logging with redaction (never log actual tokens) verbose_logger.debug("Databricks request headers: %s", self.redact_headers_for_logging(headers)) - endpoint_path: Final = "" if custom_endpoint is True else DATABRICKS_ENDPOINT_PATHS[endpoint_type] - return f"{api_base}{endpoint_path}", headers + if custom_endpoint is True: + return api_base, headers + return join_api_base(api_base, DATABRICKS_ENDPOINT_PATHS[endpoint_type]), headers diff --git a/litellm/llms/databricks/decisions/transformation.py b/litellm/llms/databricks/decisions/transformation.py index c2ac1f220b8..d1698fdf5e2 100644 --- a/litellm/llms/databricks/decisions/transformation.py +++ b/litellm/llms/databricks/decisions/transformation.py @@ -6,6 +6,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig from litellm.types.decisions import DecisionsIRRequest, UnsupportedDecisionsRequest @@ -84,8 +85,8 @@ class DatabricksDecisionsConfig(BaseDecisionsConfig): def get_complete_url(self, api_base: str, model: str) -> str: if is_ai_decide_model(model): - return f"{databricks_workspace_host(api_base)}{_AI_DECIDE_PATH}" - return f"{api_base.rstrip('/')}/{validate_databricks_decisions_model(model)}/invocations" + return join_api_base(databricks_workspace_host(api_base), _AI_DECIDE_PATH) + return join_api_base(api_base, f"{validate_databricks_decisions_model(model)}/invocations") def transform_decisions_request( self, diff --git a/litellm/llms/databricks/embed/transformation.py b/litellm/llms/databricks/embed/transformation.py index ddaf2bbb621..0feb2aea5a0 100644 --- a/litellm/llms/databricks/embed/transformation.py +++ b/litellm/llms/databricks/embed/transformation.py @@ -3,6 +3,7 @@ Translates from OpenAI's `/v1/embeddings` to Databricks' `/embeddings` """ import types +from collections.abc import Mapping from typing import Final @@ -16,16 +17,17 @@ class DatabricksEmbeddingConfig: ) def __init__(self, instruction: str | None = None) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/databricks/messages/transformation.py b/litellm/llms/databricks/messages/transformation.py index cb8d2b652a2..f5ccbd92fb4 100644 --- a/litellm/llms/databricks/messages/transformation.py +++ b/litellm/llms/databricks/messages/transformation.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator, Mapping from typing import Final import httpx +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items @@ -25,6 +26,7 @@ DROP_CONTAINER_WARNING: Final = ( "and the chat completions bridge it replaces never forwarded it." ) BEDROCK_INVOCATION_METRICS_FIELD: Final = "amazon-bedrock-invocationMetrics" +_DECODED_JSON: Final = TypeAdapter(object) def _is_unsupported_server_tool(tool: object) -> bool: @@ -57,7 +59,7 @@ def _without_unsupported_server_tools(params: Mapping[str, object]) -> dict[str, def _without_invocation_metrics(sse_line: bytes) -> bytes: if not sse_line.startswith(b"data:") or BEDROCK_INVOCATION_METRICS_FIELD.encode() not in sse_line: return sse_line - event: Final = as_str_mapping(json.loads(sse_line.removeprefix(b"data:"))) + event: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(sse_line.removeprefix(b"data:")))) if event is None: return sse_line kept_fields: Final = {key: value for key, value in event.items() if key != BEDROCK_INVOCATION_METRICS_FIELD} diff --git a/litellm/llms/databricks/responses/transformation.py b/litellm/llms/databricks/responses/transformation.py index 804c79ea79d..55c3edfefd0 100644 --- a/litellm/llms/databricks/responses/transformation.py +++ b/litellm/llms/databricks/responses/transformation.py @@ -10,6 +10,7 @@ Reference: https://docs.databricks.com/aws/en/machine-learning/foundation-model- import os from typing import TYPE_CHECKING, Any, Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.databricks.common_utils import DatabricksBase from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import ResponseInputParam @@ -70,8 +71,7 @@ class DatabricksResponsesAPIConfig(DatabricksBase, OpenAIResponsesAPIConfig): ) -> str: api_base = api_base or os.getenv("DATABRICKS_API_BASE") api_base = self._get_api_base(api_base) - api_base = api_base.rstrip("/") - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def transform_responses_api_request( self, diff --git a/litellm/llms/deepagents/harness/sandbox_backend.py b/litellm/llms/deepagents/harness/sandbox_backend.py index d39690ccdd1..8ad90faa2b4 100644 --- a/litellm/llms/deepagents/harness/sandbox_backend.py +++ b/litellm/llms/deepagents/harness/sandbox_backend.py @@ -22,7 +22,7 @@ import shlex import uuid from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import Any, Final, TypeVar +from typing import Final, TypeVar from deepagents.backends.protocol import ( DeleteResult, @@ -171,7 +171,7 @@ class SandboxBackend(SandboxBackendProtocol): # pyright: ignore[reportUntypedBa # -- sync bridge (deepagents only calls these outside the event loop) --- - def _sync(self, coro: Coroutine[Any, Any, T]) -> T: + def _sync(self, coro: Coroutine[object, object, T]) -> T: try: running = asyncio.get_running_loop() except RuntimeError: diff --git a/litellm/llms/deepagents/harness/transformation.py b/litellm/llms/deepagents/harness/transformation.py index cb80e3e12f7..54a6a5573ac 100644 --- a/litellm/llms/deepagents/harness/transformation.py +++ b/litellm/llms/deepagents/harness/transformation.py @@ -11,11 +11,13 @@ from __future__ import annotations import itertools import json -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol +from pydantic import TypeAdapter + from litellm.harness.options import DeepAgentsOptions from litellm.harness.types import ( Capabilities, @@ -41,6 +43,7 @@ DEEPAGENTS_DEFAULT_RECURSION_LIMIT: Final = 1000 _MODEL_NODE: Final = "model" # Only these graph nodes produce new messages; middleware hooks may re-emit history. _EVENT_NODES: Final = frozenset({"model", "tools"}) +_TOOL_CALLS: Final = TypeAdapter(Iterable[Mapping[str, object]]) NORMALIZED_TO_NATIVE: Final[Mapping[str, str]] = MappingProxyType( { @@ -208,7 +211,7 @@ def update_events( def _message_events(message: object, skip_tools: frozenset[str]) -> tuple[Event, ...]: kind = getattr(message, "type", None) if kind == "ai": - calls = getattr(message, "tool_calls", None) or () + calls: Final = _TOOL_CALLS.validate_python(getattr(message, "tool_calls", None) or ()) return tuple(tool_call_event(call) for call in calls if call.get("name") not in skip_tools) if kind == "tool" and getattr(message, "name", None) not in skip_tools: status: Final[object] = getattr(message, "status", None) diff --git a/litellm/llms/deepgram/audio_transcription/transformation.py b/litellm/llms/deepgram/audio_transcription/transformation.py index a3812dac8ef..1c8a056116f 100644 --- a/litellm/llms/deepgram/audio_transcription/transformation.py +++ b/litellm/llms/deepgram/audio_transcription/transformation.py @@ -2,14 +2,16 @@ Translates from OpenAI's `/v1/audio/transcriptions` to Deepgram's `/v1/listen` """ -from collections.abc import Iterable, Mapping -from typing import Final +from collections.abc import Iterable, Mapping, Sequence +from typing import Final, Protocol, cast from urllib.parse import urlencode from httpx import Headers, Response -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, InstanceOf, TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -26,6 +28,31 @@ from ..common_utils import DeepgramException _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_RESPONSE_BODY: Final = TypeAdapter(InstanceOf[dict[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_WORD_LIST: Final = TypeAdapter(InstanceOf[list[object]], config=ConfigDict(hide_input_in_errors=True)) + + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + +class _DeepgramAlternative(TypedDict): + transcript: ReadOnly[str] + words: ReadOnly[Sequence[Mapping[str, object]]] + paragraphs: ReadOnly[NotRequired[Mapping[str, str]]] + + +class _DeepgramChannel(TypedDict): + alternatives: ReadOnly[Sequence[_DeepgramAlternative]] + + +class _DeepgramResponse(TypedDict): + results: ReadOnly[Mapping[str, Sequence[_DeepgramChannel]]] + metadata: ReadOnly[Mapping[str, float]] class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @@ -81,7 +108,10 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): Transforms the raw response from Deepgram to the TranscriptionResponse format """ try: - response_json: Final = raw_response.json() + body: Final = _json_body(raw_response) + response_json: Final = cast( # cast-ok: the body passes through unvalidated, as before it was typed + _DeepgramResponse, body + ) # Get the first alternative from the first channel first_channel: Final = response_json["results"]["channels"][0] @@ -101,7 +131,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): text = first_alternative["paragraphs"]["transcript"] else: # Diarization without paragraphs: reconstruct from words - text = self._reconstruct_diarized_transcript(first_alternative["words"]) + text = self._reconstruct_diarized_transcript(_WORD_LIST.validate_python(first_alternative["words"])) # Create TranscriptionResponse object response: Final = TranscriptionResponse(text=text) @@ -123,7 +153,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ] # Store full response in hidden params - response.hidden_params = response_json + response.hidden_params = _RESPONSE_BODY.validate_python(body) return response @@ -179,7 +209,6 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> str: if api_base is None: api_base = get_secret_str("DEEPGRAM_API_BASE") or "https://api.deepgram.com/v1" - api_base = api_base.rstrip("/") # Remove trailing slash if present # Build query parameters including the model all_query_params: Final = {"model": model} @@ -189,7 +218,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): all_query_params.update(additional_params) # Construct URL with proper query string encoding - base_url: Final = f"{api_base}/listen" + base_url: Final = join_api_base(api_base, "listen") query_string: Final = urlencode(all_query_params) url: Final = f"{base_url}?{query_string}" diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index 9b071ac8321..29998895bbc 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -8,6 +8,7 @@ import httpx import litellm from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import LlmProviders @@ -50,7 +51,7 @@ def _first_occurrences(query_string: str) -> httpx.QueryParams: def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str: - listen_url: Final = httpx.URL(f"{(api_base or DEEPGRAM_DEFAULT_API_BASE).rstrip('/')}/listen") + listen_url: Final = httpx.URL(join_api_base(api_base or DEEPGRAM_DEFAULT_API_BASE, "listen")) websocket_url: Final = listen_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(listen_url.scheme, listen_url.scheme)) params: Final = _first_occurrences(query_string) query: Final = params if params.get("model") else params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL) diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 5196c44d504..62a780396a1 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -6,6 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Final, Protocol import httpx +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._uuid import uuid @@ -24,6 +25,8 @@ from litellm.types.rerank import ( RerankTokens, ) +_DECODED_JSON: Final = TypeAdapter(object) + class _DeepinfraInferenceStatus(TypedDict, total=False): """The ``inference_status`` block of a DeepInfra rerank response.""" @@ -251,7 +254,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): try: error_data = error_message if isinstance(error_message, str): - error_data = json.loads(error_message) + error_data = _DECODED_JSON.validate_python(json.loads(error_message)) if isinstance(error_data, dict): # Check for {"detail": {"error": "..."}} detail: Final = error_data.get("detail") diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index d3319856d29..ecaa0a2ef16 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -10,12 +10,15 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, extract_search_results_text, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_reasoning, supports_vision from ...openai.chat.gpt_transformation import OpenAIGPTConfig +SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "xhigh", "max", "ultra") + class DeepSeekChatConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> list: @@ -32,7 +35,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): optional_params: dict, model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: """ Map OpenAI params to DeepSeek params. @@ -59,6 +62,13 @@ class DeepSeekChatConfig(OpenAIGPTConfig): elif reasoning_effort is not None: optional_params["thinking"] = {"type": "disabled" if reasoning_effort == "none" else "enabled"} + if ( + isinstance(reasoning_effort, str) + and reasoning_effort in SUPPORTED_REASONING_EFFORTS + and optional_params["thinking"] == {"type": "enabled"} + ): + return {**optional_params, "reasoning_effort": reasoning_effort} + return optional_params def _fill_reasoning_content(self, messages: list[AllMessageValues]) -> list[AllMessageValues]: @@ -363,9 +373,4 @@ class DeepSeekChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default DeepSeek /chat/completions endpoint. """ - clean_api_base: Final = (api_base or "https://api.deepseek.com/beta").rstrip("/") - - if not clean_api_base.endswith("/chat/completions"): - return f"{clean_api_base}/chat/completions" - - return clean_api_base + return ensure_api_route(api_base or "https://api.deepseek.com/beta", "chat/completions") diff --git a/litellm/llms/deepseek/messages/transformation.py b/litellm/llms/deepseek/messages/transformation.py index 6a6b57131ad..3b2dd5af260 100644 --- a/litellm/llms/deepseek/messages/transformation.py +++ b/litellm/llms/deepseek/messages/transformation.py @@ -2,16 +2,37 @@ DeepSeek Anthropic-compatible messages transformation config. """ -from collections.abc import Sequence -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final + +from pydantic import TypeAdapter import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams +_TOOL_LIST: Final = TypeAdapter(list[object]) +_TOOL_FIELDS: Final = TypeAdapter(dict[object, object]) + + +def _custom_tool_without_type(tool: object) -> Mapping[object, object] | None: + if not isinstance(tool, dict): + return None + fields: Final = _TOOL_FIELDS.validate_python(tool) + if fields.get("type") == "custom": + fields.pop("type", None) + return fields + return None + + +def _sanitized_tool(tool: object) -> object: + stripped: Final = _custom_tool_without_type(tool) + return tool if stripped is None else stripped + class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): """ @@ -90,24 +111,16 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): base_url = base_url.removesuffix("/beta") if not base_url.endswith("/anthropic") and "/anthropic/" not in base_url: - base_url = f"{base_url}/anthropic" + base_url = join_api_base(base_url, "anthropic") - return f"{base_url}/v1/messages" + return join_api_base(base_url, "v1/messages") @staticmethod - def _sanitize_tools_for_deepseek(tools: Any) -> object: + def _sanitize_tools_for_deepseek(tools: object) -> object: if not isinstance(tools, list): return tools - sanitized_tools: Final = [] - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "custom": - sanitized_tool = dict(tool) - sanitized_tool.pop("type", None) - sanitized_tools.append(sanitized_tool) - else: - sanitized_tools.append(tool) - return sanitized_tools + return [_sanitized_tool(tool) for tool in _TOOL_LIST.validate_python(tools)] def transform_anthropic_messages_request( self, diff --git a/litellm/llms/deprecated_providers/palm.py b/litellm/llms/deprecated_providers/palm.py index 785cffa48ea..b4a32630c42 100644 --- a/litellm/llms/deprecated_providers/palm.py +++ b/litellm/llms/deprecated_providers/palm.py @@ -1,4 +1,5 @@ import types +from collections.abc import Mapping from typing import Final @@ -41,16 +42,17 @@ class PalmConfig: top_p: float | None = None, max_output_tokens: int | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/docker_model_runner/chat/transformation.py b/litellm/llms/docker_model_runner/chat/transformation.py index adb40311efd..68b63812248 100644 --- a/litellm/llms/docker_model_runner/chat/transformation.py +++ b/litellm/llms/docker_model_runner/chat/transformation.py @@ -10,6 +10,7 @@ from typing import Final, Literal, overload from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -104,12 +105,9 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): if not api_base: api_base = "http://localhost:22088/engines/llama.cpp" - # Remove trailing slashes from api_base - api_base = api_base.rstrip("/") - # Build the URL: {api_base}/v1/chat/completions # api_base is expected to already contain the engine path - complete_url: Final = f"{api_base}/v1/chat/completions" + complete_url: Final = join_api_base(api_base, "v1/chat/completions") return complete_url diff --git a/litellm/llms/duckduckgo/search/transformation.py b/litellm/llms/duckduckgo/search/transformation.py index 1816a0d24e4..36893b4bd27 100644 --- a/litellm/llms/duckduckgo/search/transformation.py +++ b/litellm/llms/duckduckgo/search/transformation.py @@ -12,6 +12,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -91,7 +92,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): if data and isinstance(data, dict) and "_duckduckgo_params" in data: params: Final = data["_duckduckgo_params"] query_string: Final = urlencode(params, doseq=True) - return f"{api_base}/?{query_string}" + return f"{join_api_base(api_base, '')}?{query_string}" return api_base diff --git a/litellm/llms/e2b/sandbox/transformation.py b/litellm/llms/e2b/sandbox/transformation.py index dbf3996248a..b5f42632eef 100644 --- a/litellm/llms/e2b/sandbox/transformation.py +++ b/litellm/llms/e2b/sandbox/transformation.py @@ -15,6 +15,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.sandbox.transformation import ( SANDBOX_MAX_OUTPUT_BYTES, BaseSandboxConfig, @@ -78,7 +79,7 @@ class E2BSandboxConfig(BaseSandboxConfig): body["metadata"] = metadata response: Final = await self._http(client).post( - url=f"{base}/sandboxes", + url=join_api_base(base, "sandboxes"), headers={"X-API-Key": key, "Content-Type": "application/json"}, json=body, ) @@ -151,10 +152,10 @@ class E2BSandboxConfig(BaseSandboxConfig): dict[str, object], getattr(handle, HIDDEN_PARAMS_ATTR) ) key: Final = api_key or hidden_params.get("api_key") or self.validate_environment() - base: Final = api_base or hidden_params.get("api_base") or E2B_API_BASE + base: Final = str(api_base or hidden_params.get("api_base") or E2B_API_BASE) try: response: Final = await self._http(client).delete( - url=f"{base}/sandboxes/{handle.id}", + url=join_api_base(base, f"sandboxes/{handle.id}"), headers={"X-API-Key": key}, ) except httpx.HTTPStatusError as e: diff --git a/litellm/llms/edenai/chat/transformation.py b/litellm/llms/edenai/chat/transformation.py index ae5da5d1cfd..480b382076d 100644 --- a/litellm/llms/edenai/chat/transformation.py +++ b/litellm/llms/edenai/chat/transformation.py @@ -23,7 +23,7 @@ from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, ModelResponseStream, Usage -from ..common_utils import EdenAIException, reported_cost, resolve_api_base, resolve_api_key +from ..common_utils import EdenAIException, endpoint_url, reported_cost, resolve_api_base, resolve_api_key if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -138,7 +138,7 @@ class EdenAIChatConfig(OpenAIGPTConfig): def get_models( self, api_key: str | None = None, api_base: str | None = None ) -> list[str]: # mutable-ok: inherited contract - response: Final = litellm.module_level_client.get(url=f"{self.get_api_base(api_base)}/models") + response: Final = litellm.module_level_client.get(url=endpoint_url(api_base, "models")) if not response.is_success: raise EdenAIException(status_code=response.status_code, message=response.text, headers=response.headers) catalog: Final = _EdenAIModelCatalog.model_validate(response.json()) diff --git a/litellm/llms/edenai/common_utils.py b/litellm/llms/edenai/common_utils.py index 3216f798fbe..60c83f7ea75 100644 --- a/litellm/llms/edenai/common_utils.py +++ b/litellm/llms/edenai/common_utils.py @@ -11,6 +11,7 @@ from pydantic import AliasChoices, Field, ValidationError import litellm from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.base import LiteLLMBaseModel @@ -74,7 +75,7 @@ def json_headers( def endpoint_url(api_base: str | None, path: str) -> str: - return f"{resolve_api_base(api_base).rstrip('/')}/{path}" + return join_api_base(resolve_api_base(api_base), path) def pick(params: Mapping[str, object], keys: Container[str]) -> Mapping[str, object]: diff --git a/litellm/llms/edenai/embedding/transformation.py b/litellm/llms/edenai/embedding/transformation.py index c79a6839434..13fa7142916 100644 --- a/litellm/llms/edenai/embedding/transformation.py +++ b/litellm/llms/edenai/embedding/transformation.py @@ -5,7 +5,7 @@ per-request cost at the top level of the body. Docs: https://www.edenai.co/docs/v3/llms/embeddings """ -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol, cast import httpx @@ -24,6 +24,14 @@ if TYPE_CHECKING: _SUPPORTED_PARAMS: Final = ("dimensions", "encoding_format", "user") +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class EdenAIEmbeddingConfig(BaseEmbeddingConfig): def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract return list(_SUPPORTED_PARAMS) @@ -80,10 +88,14 @@ class EdenAIEmbeddingConfig(BaseEmbeddingConfig): optional_params: dict[str, object], # mutable-ok: inherited contract litellm_params: dict[str, object], # mutable-ok: inherited contract ) -> EmbeddingResponse: - body: Final = raw_response.json() + body: Final = _json_body(raw_response) logging_obj.post_call(original_response=body) response: Final[EmbeddingResponse] = convert_to_model_response_object( - response_object=body, model_response_object=model_response, response_type="embedding" + response_object=cast( # cast-ok: the body passes through unvalidated, as before it was typed + dict[str, object], body + ), + model_response_object=model_response, + response_type="embedding", ) set_response_cost_in_hidden_params(response, reported_cost(body)) return response diff --git a/litellm/llms/edenai/image_generation/transformation.py b/litellm/llms/edenai/image_generation/transformation.py index e2a0b8684f6..a0c4e1b0458 100644 --- a/litellm/llms/edenai/image_generation/transformation.py +++ b/litellm/llms/edenai/image_generation/transformation.py @@ -5,7 +5,7 @@ for every image model in the catalog with the real per-request cost at the top l Docs: https://www.edenai.co/docs/v3/llms/image-generation """ -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol, cast import httpx @@ -36,6 +36,14 @@ _SUPPORTED_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = ( ) +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class EdenAIImageGenerationConfig(BaseImageGenerationConfig): def get_supported_openai_params( self, model: str @@ -97,10 +105,14 @@ class EdenAIImageGenerationConfig(BaseImageGenerationConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: - body: Final = raw_response.json() + body: Final = _json_body(raw_response) logging_obj.post_call(original_response=body) response: Final[ImageResponse] = convert_to_model_response_object( - response_object=body, model_response_object=model_response, response_type="image_generation" + response_object=cast( # cast-ok: the body passes through unvalidated, as before it was typed + dict[str, object], body + ), + model_response_object=model_response, + response_type="image_generation", ) set_response_cost_in_hidden_params(response, reported_cost(body)) return response diff --git a/litellm/llms/elevenlabs/audio_transcription/transformation.py b/litellm/llms/elevenlabs/audio_transcription/transformation.py index 04d2717d64f..3dc42f35d3c 100644 --- a/litellm/llms/elevenlabs/audio_transcription/transformation.py +++ b/litellm/llms/elevenlabs/audio_transcription/transformation.py @@ -3,13 +3,14 @@ Translates from OpenAI's `/v1/audio/transcriptions` to ElevenLabs's `/v1/speech- """ from collections.abc import Iterable, Mapping -from typing import Final +from typing import Final, Protocol from httpx import Headers, Response from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -26,6 +27,15 @@ from ..common_utils import ElevenLabsException _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_HIDDEN_PARAMS: Final = TypeAdapter(dict[str, object]) + + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @@ -119,7 +129,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): Transforms the raw response from ElevenLabs to the TranscriptionResponse format """ try: - response_json: Final = raw_response.json() + response_json: Final = _json_body(raw_response) response_object: Final = _JSON_OBJECT.validate_python(response_json) # Extract the main transcript text @@ -134,20 +144,18 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # Map ElevenLabs words to OpenAI format if "words" in response_object: - response["words"] = [] - for word_data in _JSON_OBJECTS.validate_python(response_object["words"]): - # Only include actual words, skip spacing and audio events - if word_data.get("type") == "word": - response["words"].append( - { - "word": word_data.get("text", ""), - "start": word_data.get("start", 0), - "end": word_data.get("end", 0), - } - ) + response["words"] = [ + { + "word": word_data.get("text", ""), + "start": word_data.get("start", 0), + "end": word_data.get("end", 0), + } + for word_data in _JSON_OBJECTS.validate_python(response_object["words"]) + if word_data.get("type") == "word" + ] # Store full response in hidden params - response.hidden_params = response_json + response.hidden_params = _HIDDEN_PARAMS.validate_python(response_json) return response @@ -165,12 +173,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> str: if api_base is None: api_base = get_secret_str("ELEVENLABS_API_BASE") or "https://api.elevenlabs.io" - api_base = api_base.rstrip("/") # Remove trailing slash if present - - # ElevenLabs speech-to-text endpoint - url: Final = f"{api_base}/v1/speech-to-text" - - return url + return join_api_base(api_base, "v1/speech-to-text") def validate_environment( self, diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index eb93543df46..23ccf41c9da 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -11,7 +11,7 @@ import httpx from httpx import Headers import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, @@ -56,9 +56,9 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): # Response format mappings from OpenAI to ElevenLabs FORMAT_MAPPINGS = { "mp3": "mp3_44100_128", - "pcm": "pcm_44100", + "pcm": "pcm_24000", "opus": "opus_48000_128", - # ElevenLabs does not support WAV, AAC, or FLAC formats. + "wav": "wav_24000", } ELEVENLABS_QUERY_PARAMS_KEY = "__elevenlabs_query_params__" @@ -297,15 +297,14 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): """ Construct the ElevenLabs endpoint URL, including path voice_id and query params. """ - base_url = api_base or get_secret_str("ELEVENLABS_API_BASE") or self.TTS_BASE_URL - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("ELEVENLABS_API_BASE") or self.TTS_BASE_URL voice_id: Final = litellm_params.get(self.ELEVENLABS_VOICE_ID_KEY) if not isinstance(voice_id, str) or not voice_id.strip(): raise ValueError("ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`.") encoded_voice_id: Final = encode_url_path_segment(voice_id, field_name="voice_id") - url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}" + url = join_api_base(base_url, f"{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}") query_params: Final = litellm_params.get(self.ELEVENLABS_QUERY_PARAMS_KEY, {}) if query_params: diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 89073409599..73e341d0033 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -93,11 +94,7 @@ class ExaAISearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("EXA_API_BASE") or self.EXA_AI_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("EXA_API_BASE") or self.EXA_AI_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py index 6d58caeb384..344d9d3c75b 100644 --- a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py +++ b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py @@ -4,6 +4,7 @@ from typing import Final from httpx._types import RequestFiles +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams @@ -45,8 +46,8 @@ class FalAIFluxLoraDepthEditConfig(FalAIImageEditConfig): api_base: str | None, litellm_params: dict, # mutable-ok: base class contract ) -> str: - base_url: Final = (api_base or get_secret_str("FAL_AI_API_BASE") or DEFAULT_BASE_URL).rstrip("/") - return f"{base_url}/{FLUX_LORA_DEPTH_ENDPOINT}" + base_url: Final = api_base or get_secret_str("FAL_AI_API_BASE") or DEFAULT_BASE_URL + return join_api_base(base_url, FLUX_LORA_DEPTH_ENDPOINT) def transform_image_edit_request( self, diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py index 0fdd89aec91..fc215e3d94e 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py @@ -19,6 +19,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class FalAIFluxProV11UltraConfig(FalAIBaseConfig): @@ -222,7 +223,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): } """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py index 038a06c5749..26b257f5c89 100644 --- a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py @@ -18,6 +18,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class FalAIRecraftV3Config(FalAIBaseConfig): @@ -195,7 +196,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): } """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py index 30d51cc8084..fdd552efd73 100644 --- a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py +++ b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py @@ -19,6 +19,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class FalAIStableDiffusionConfig(FalAIBaseConfig): @@ -236,7 +237,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): } """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index 0bb5dada071..79913d5c651 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -85,12 +86,10 @@ class FalAIBaseConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - if self.IMAGE_GENERATION_ENDPOINT: - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" - return complete_url + base_url: Final = api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL + if not self.IMAGE_GENERATION_ENDPOINT: + return base_url.rstrip("/") + return join_api_base(base_url, self.IMAGE_GENERATION_ENDPOINT) def validate_environment( self, diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 3136763e1be..86aa6388ff0 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -10,7 +10,7 @@ import httpx from httpx._types import FileContent, RequestFiles from pydantic import TypeAdapter -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( @@ -445,7 +445,7 @@ class FalAIVideoConfig(BaseVideoConfig): request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( - f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}/status", + join_api_base(api_base, f"{_queue_request_base_path(model_id)}/requests/{encoded_request_id}/status"), {}, ) @@ -541,7 +541,7 @@ class FalAIVideoConfig(BaseVideoConfig): request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( - f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}", + join_api_base(api_base, f"{_queue_request_base_path(model_id)}/requests/{encoded_request_id}"), {}, ) diff --git a/litellm/llms/fastcrw/search/transformation.py b/litellm/llms/fastcrw/search/transformation.py index 8f6f8d87300..27db6a95170 100644 --- a/litellm/llms/fastcrw/search/transformation.py +++ b/litellm/llms/fastcrw/search/transformation.py @@ -15,6 +15,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -83,11 +84,7 @@ class FastCRWSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("CRW_API_BASE") or self.FASTCRW_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("CRW_API_BASE") or self.FASTCRW_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/featherless_ai/chat/transformation.py b/litellm/llms/featherless_ai/chat/transformation.py index a68261c58e9..4bf0486f723 100644 --- a/litellm/llms/featherless_ai/chat/transformation.py +++ b/litellm/llms/featherless_ai/chat/transformation.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final import litellm @@ -42,7 +44,7 @@ class FeatherlessAIConfig(OpenAIGPTConfig): tool_choice: str | None = None, tools: list | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/firecrawl/search/transformation.py b/litellm/llms/firecrawl/search/transformation.py index ee99ab24413..6d1e2b22918 100644 --- a/litellm/llms/firecrawl/search/transformation.py +++ b/litellm/llms/firecrawl/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -85,11 +86,7 @@ class FirecrawlSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("FIRECRAWL_API_BASE") or self.FIRECRAWL_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("FIRECRAWL_API_BASE") or self.FIRECRAWL_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 14ae59d4f85..fa5794d71b3 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_response_only_fields, unpack_legacy_defs, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -826,7 +827,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): base = api_base.rstrip("/") base = base.removesuffix("/v1") response: Final = litellm.module_level_client.get( - url=f"{base}/v1/accounts/{account_id}/models", + url=join_api_base(base, f"v1/accounts/{account_id}/models"), headers={"Authorization": f"Bearer {api_key}"}, ) diff --git a/litellm/llms/fireworks_ai/decisions/transformation.py b/litellm/llms/fireworks_ai/decisions/transformation.py index 87d402902be..bba4835060b 100644 --- a/litellm/llms/fireworks_ai/decisions/transformation.py +++ b/litellm/llms/fireworks_ai/decisions/transformation.py @@ -14,6 +14,8 @@ class FireworksAIDecisionsConfig(BaseDecisionsConfig): def parse_response(self, payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse: response: Final = parse_systemone_response(self.unwrap_response(payload), request) + if response.usage is None: + return response cached: Final = response.usage.extra.get("cached_input_tokens") if isinstance(cached, bool) or not isinstance(cached, int): return response diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 4eff64e351f..84c93974cd7 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -12,6 +12,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.fireworks_ai.common_utils import FireworksAIMixin from litellm.types.llms.base import LiteLLMBaseModel @@ -43,6 +44,7 @@ class _FireworksAIResultFields(LiteLLMBaseModel): _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _STR: Final = TypeAdapter(str) +_DECODED_JSON: Final = TypeAdapter(object) class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): @@ -61,9 +63,9 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): api_base = api_base.rstrip("/") if not api_base.endswith("/rerank"): if api_base.endswith("/v1") or api_base.endswith("/inference/v1"): - api_base = f"{api_base}/rerank" + api_base = join_api_base(api_base, "rerank") else: - api_base = f"{api_base}/inference/v1/rerank" + api_base = join_api_base(api_base, "inference/v1/rerank") return api_base return "https://api.fireworks.ai/inference/v1/rerank" @@ -199,7 +201,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): Transform Fireworks AI rerank response to LiteLLM RerankResponse format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Failed to parse response: {e}", diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index 30eaa5a7ca2..db234c09930 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -6,6 +6,7 @@ from urllib.parse import unquote import httpx from openai.types.responses import EasyInputMessageParam, ResponseInputContentParam, ResponseInputItemParam +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.fireworks_ai.common_utils import ( FIREWORKS_FORWARD_USER_ID_PARAM, get_fireworks_forwarded_user_id, @@ -150,8 +151,8 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return dict(pinned) def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: - base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/") - return f"{base}/responses" + base: Final = api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE + return join_api_base(base, "responses") def transform_responses_api_request( self, diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index 6eac3ac79cd..680db52b283 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -10,12 +10,17 @@ from collections.abc import Callable from typing import Final, Protocol from urllib.parse import urlsplit +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict, Unpack import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.types.llms.openai import AllMessageValues +_DECODED_JSON: Final = TypeAdapter(object) + class _OpenAIGPTConfigOptions(TypedDict, total=False): """The sampling defaults ``OpenAIGPTConfig.__init__`` accepts and stashes on the class.""" @@ -140,7 +145,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): project = self._validate_path_id(project, "vertex_project", model) location = self._validate_path_id(location, "vertex_location", model) - return f"{api_base}/v1/projects/{project}/locations/{location}/chat/completions" + return join_api_base(api_base, f"v1/projects/{project}/locations/{location}/chat/completions") def _read_env_bool(self, val: bool | str | None, env_var: str, default: bool = True) -> bool | str: def _parse(s: str) -> bool | str: @@ -202,10 +207,10 @@ class GDCGeminiConfig(OpenAILikeChatConfig): import google.auth try: - json_obj: Final = json.loads(api_key) + json_obj: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(api_key))) except json.JSONDecodeError: return None, False - if not isinstance(json_obj, dict) or json_obj.get("type") != self._GDCH_CREDENTIAL_TYPE: + if json_obj is None or json_obj.get("type") != self._GDCH_CREDENTIAL_TYPE: raise ValueError( "GDC only accepts a GDCH service account credential as a JSON api_key " '(expected "type": "gdch_service_account"). Other Google credential types are ' diff --git a/litellm/llms/gemini/agents/transformation.py b/litellm/llms/gemini/agents/transformation.py index cfdb55f2048..f7352da5fb6 100644 --- a/litellm/llms/gemini/agents/transformation.py +++ b/litellm/llms/gemini/agents/transformation.py @@ -10,11 +10,13 @@ Proxies the Gemini v1beta Agents API: """ from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, Protocol import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo from litellm.types.agents import ( @@ -28,6 +30,15 @@ from litellm.types.agents import ( # create-agent body verbatim. _GEMINI_AGENT_BODY_KEYS: Final = ("base_agent", "instructions", "base_environment") + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + # LiteLLM-internal keys that must never be forwarded to Gemini. _LITELLM_INTERNAL_KEYS: Final = frozenset( { @@ -67,7 +78,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): return "v1beta" def _base_url(self, api_base: str | None) -> str: - return f"{GeminiModelInfo.get_api_base(api_base)}/{self.api_version}" + return join_api_base(GeminiModelInfo.get_api_base(api_base), self.api_version) # ------------------------------------------------------------------ # # Shared helpers # @@ -90,7 +101,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> str: - return f"{self._base_url(api_base)}/agents" + return join_api_base(self._base_url(api_base), "agents") def validate_environment( self, @@ -155,18 +166,19 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): """ self._raise_for_status(raw_response) try: - data: dict[str, Any] = raw_response.json() + decoded: object = _json_body(raw_response) except Exception: verbose_logger.warning( "GeminiAgentsConfig: non-JSON create response (status=%d).", raw_response.status_code, ) - data = {"id": name} + decoded = {"id": name} + data: Final = dict(decoded_json_object(decoded)) # Gemini uses "id" as the identifier; normalise to both fields. data.setdefault("id", name) data.setdefault("name", data["id"]) verbose_logger.debug("GeminiAgentsConfig create response: %s", data) - return AgentCreateResponse(**data) + return AgentCreateResponse.model_validate(data) # ------------------------------------------------------------------ # # LIST # @@ -177,7 +189,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, object]]: - url: Final = f"{self._base_url(api_base)}/agents" + url: Final = join_api_base(self._base_url(api_base), "agents") params: Final[dict[str, object]] = {} if litellm_params.get("page_size"): params["pageSize"] = litellm_params["page_size"] @@ -191,13 +203,16 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): ) -> AgentListResponse: self._raise_for_status(raw_response) try: - data = raw_response.json() + decoded: object = _json_body(raw_response) except Exception: - data = {} + decoded = {} + data: Final = decoded_json_object(decoded) verbose_logger.debug("GeminiAgentsConfig list response: %s", data) - return AgentListResponse( - agents=data.get("agents", []), - next_page_token=data.get("nextPageToken"), + return AgentListResponse.model_validate( + { + "agents": data.get("agents", []), + "next_page_token": data.get("nextPageToken"), + } ) # ------------------------------------------------------------------ # @@ -210,7 +225,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, object]]: - url: Final = f"{self._base_url(api_base)}/agents/{name}" + url: Final = join_api_base(self._base_url(api_base), f"agents/{name}") return url, {} def transform_get_response( @@ -221,13 +236,14 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): """Same shape as create response — Gemini returns "id" as identifier.""" self._raise_for_status(raw_response) try: - data = raw_response.json() + decoded: object = _json_body(raw_response) except Exception: - data = {"id": name} + decoded = {"id": name} + data: Final = dict(decoded_json_object(decoded)) data.setdefault("id", name) data.setdefault("name", data["id"]) verbose_logger.debug("GeminiAgentsConfig get response: %s", data) - return AgentCreateResponse(**data) + return AgentCreateResponse.model_validate(data) # ------------------------------------------------------------------ # # DELETE # @@ -239,7 +255,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> str: - return f"{self._base_url(api_base)}/agents/{name}" + return join_api_base(self._base_url(api_base), f"agents/{name}") def transform_delete_response( self, @@ -265,7 +281,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, object]]: - url: Final = f"{self._base_url(api_base)}/agents/{name}/versions" + url: Final = join_api_base(self._base_url(api_base), f"agents/{name}/versions") params: Final[dict[str, object]] = {} if litellm_params.get("page_size"): params["pageSize"] = litellm_params["page_size"] @@ -284,11 +300,14 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): """ self._raise_for_status(raw_response) try: - data = raw_response.json() + decoded: object = _json_body(raw_response) except Exception: - data = {} + decoded = {} + data: Final = decoded_json_object(decoded) verbose_logger.debug("GeminiAgentsConfig list_versions response for '%s': %s", name, data) - return AgentVersionsResponse( - agent_versions=data.get("agentVersions", []), - next_page_token=data.get("nextPageToken"), + return AgentVersionsResponse.model_validate( + { + "agent_versions": data.get("agentVersions", []), + "next_page_token": data.get("nextPageToken"), + } ) diff --git a/litellm/llms/gemini/audio_transcription/transformation.py b/litellm/llms/gemini/audio_transcription/transformation.py index 48a345a4940..3457b22faa4 100644 --- a/litellm/llms/gemini/audio_transcription/transformation.py +++ b/litellm/llms/gemini/audio_transcription/transformation.py @@ -3,12 +3,14 @@ from collections.abc import Mapping, Sequence from typing import Final from httpx import Headers, Response +from pydantic import TypeAdapter from litellm.litellm_core_utils.audio_utils.subtitle_utils import SUBTITLE_RESPONSE_FORMATS from litellm.litellm_core_utils.audio_utils.utils import ( normalize_transcription_language_to_bcp47, process_audio_file, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -35,6 +37,7 @@ from litellm.types.utils import ( INTERACTIONS_API_REVISION: Final = "2026-05-20" WORD_INFO_ANNOTATION_TYPE: Final = "word_info" +_DECODED_JSON: Final = TypeAdapter(object) class GeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @@ -104,8 +107,7 @@ class GeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: - resolved_api_base: Final = GeminiModelInfo.get_api_base(api_base) - return f"{resolved_api_base}/v1beta/interactions" + return join_api_base(GeminiModelInfo.get_api_base(api_base), "v1beta/interactions") def transform_audio_transcription_request( self, @@ -132,7 +134,7 @@ class GeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig): raw_response: Response, ) -> TranscriptionResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except ValueError: raise GeminiError( status_code=raw_response.status_code, diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 8ee7a463828..9634c580858 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -2,15 +2,16 @@ import base64 import datetime import json import math -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, cast import httpx from pydantic import JsonValue, TypeAdapter import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -377,7 +378,7 @@ class GeminiModelInfo(BaseLLMModelInfo): return "v1beta" @staticmethod - def get_api_base(api_base: str | None = None) -> str | None: + def get_api_base(api_base: str | None = None) -> str: return api_base or get_secret_str("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" @staticmethod @@ -388,7 +389,7 @@ class GeminiModelInfo(BaseLLMModelInfo): def get_base_model(model: str) -> str | None: return model.replace("gemini/", "") - def process_model_name(self, models: list[dict[str, str]]) -> list[str]: + def process_model_name(self, models: Iterable[Mapping[str, str]]) -> list[str]: litellm_model_names: Final = [] for model in models: stripped_model_name = model["name"].replace("models/", "") @@ -400,13 +401,13 @@ class GeminiModelInfo(BaseLLMModelInfo): api_base = GeminiModelInfo.get_api_base(api_base) api_key = GeminiModelInfo.get_api_key(api_key) endpoint: Final = f"/{self.api_version}/models" - if api_base is None or api_key is None: + if api_key is None: raise ValueError( "GEMINI_API_BASE or GEMINI_API_KEY/GOOGLE_API_KEY is not set. Please set the environment variable, to query Gemini's `/models` endpoint." ) response: Final = litellm.module_level_client.get( - url=f"{api_base}{endpoint}", + url=join_api_base(api_base, endpoint), headers={"x-goog-api-key": api_key}, ) @@ -415,9 +416,11 @@ class GeminiModelInfo(BaseLLMModelInfo): f"Failed to fetch models from Gemini. Status code: {response.status_code}, Response: {response.json()}" ) - models: Final[list[dict[str, str]]] = response.json()["models"] + listing: Final = cast( # cast-ok: the body passes through unvalidated, as before it was typed + Mapping[str, Iterable[Mapping[str, str]]], _JSON_VALUE.validate_python(response.json()) + ) - litellm_model_names: Final = self.process_model_name(models) + litellm_model_names: Final = self.process_model_name(listing["models"]) return litellm_model_names def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index 1c9d09330a8..4ed47323d4f 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -6,6 +6,7 @@ from pydantic import ConfigDict, RootModel, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.utils import LlmProviders @@ -65,7 +66,7 @@ class GoogleAIStudioTokenCounter: Construct the URL for the Google Gen AI Studio countTokens endpoint. """ base_url: Final = api_base or "https://generativelanguage.googleapis.com" - return f"{base_url}/v1beta/models/{model}:countTokens" + return join_api_base(base_url, f"v1beta/models/{model}:countTokens") def __init__(self, client: AsyncHTTPHandler | None = None) -> None: super().__init__() diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index 2c62e04c5a3..f9a7e230ea9 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -15,7 +15,7 @@ from typing_extensions import ReadOnly, Required from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, @@ -106,8 +106,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): if not final_api_key: raise ValueError("api_key is required") - url: Final = f"{api_base}/{endpoint}" - return url + return join_api_base(api_base, endpoint) def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: return [] @@ -235,10 +234,11 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): file_part: Final = self._normalize_gemini_file_id(file_id) - api_base = self.get_api_base(litellm_params.get("api_base")) or "https://generativelanguage.googleapis.com" - api_base = api_base.rstrip("/") + api_base: Final = ( + self.get_api_base(litellm_params.get("api_base")) or "https://generativelanguage.googleapis.com" + ) - url: Final = f"{api_base}/v1beta/{file_part}" + url: Final = join_api_base(api_base, f"v1beta/{file_part}") # API key is passed via x-goog-api-key header (set in validate_environment) return url, {} @@ -345,7 +345,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): file_name: Final = self._normalize_gemini_file_id(file_id) # Construct the delete URL - url: Final = f"{api_base}/v1beta/{file_name}" + url: Final = join_api_base(api_base, f"v1beta/{file_name}") # Add API key as header (Google AI Studio uses x-goog-api-key header) params: Final[dict] = {} diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py index 44cdac11325..34007bc8fd4 100644 --- a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py +++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py @@ -15,7 +15,10 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamTransformSink, @@ -28,6 +31,7 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _EMPTY_REQUEST_DATA: Final[Mapping[str, object]] = MappingProxyType({}) +_DECODED_JSON: Final = TypeAdapter(object) def _field(container: object, name: str) -> object | None: @@ -108,12 +112,10 @@ def _write_back_texts(text_parts: Sequence[object], guardrailed_texts: Sequence[ def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None: try: - parsed: Final = json.loads(payload) + parsed: Final = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError: return None - if isinstance(parsed, dict): - return parsed - return None + return as_str_mapping(parsed) def _sse_payload_texts(sse_text: str) -> tuple[str, ...]: diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index bc3180a2129..0e29df36e80 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -3,7 +3,7 @@ Transformation for Calling Google models in their native format. """ from copy import deepcopy -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast import httpx @@ -36,6 +36,14 @@ else: from ..common_utils import get_api_key_from_env +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): """ Configuration for calling Google models in their native format. @@ -382,7 +390,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): from litellm.types.google_genai.main import GenerateContentResponse try: - response = raw_response.json() + decoded: Final = _json_body(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error transforming generate content response: {e}", @@ -391,7 +399,9 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): ) logging_obj.model_call_details["httpx_response"] = raw_response - response = self.convert_citation_sources_to_citations(response) + response: Final = self.convert_citation_sources_to_citations( + cast(dict[str, object], decoded) # cast-ok: the body passes through unvalidated, as before it was typed + ) return GenerateContentResponse(**response) diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 31d3963c70c..d48a66f1f93 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -1,12 +1,14 @@ import base64 from collections.abc import Mapping from io import BufferedReader, BytesIO -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import httpx from httpx._types import RequestFiles from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.core_helpers import decoded_json_object +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.gemini.common_utils import ( get_gemini_image_generation_config, @@ -17,6 +19,7 @@ from litellm.llms.gemini.image_usage_transformation import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.llms.vertex_ai import GenerateContentResponseBody from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( FileTypes, @@ -33,6 +36,14 @@ else: LiteLLMLoggingObj = Any +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class GeminiImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" SUPPORTED_PARAMS: list[str] = ["n", "size", "imageConfig"] @@ -79,9 +90,8 @@ class GeminiImageEditConfig(BaseImageEditConfig): api_base: str | None, litellm_params: dict, ) -> str: - base_url = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL - base_url = base_url.rstrip("/") - return f"{base_url}/models/{model}:generateContent" + base_url: Final = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, f"models/{model}:generateContent") def transform_image_edit_request( self, @@ -125,7 +135,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): ) -> ImageResponse: model_response: Final = ImageResponse() try: - response_json: Final = raw_response.json() + response_json: Final = _json_body(raw_response) except Exception as exc: raise self.get_error_class( error_message=f"Error transforming image edit response: {exc}", @@ -133,7 +143,10 @@ class GeminiImageEditConfig(BaseImageEditConfig): headers=raw_response.headers, ) - candidates: Final = response_json.get("candidates", []) + payload: Final = cast( # cast-ok: the body passes through unvalidated, as before it was typed + GenerateContentResponseBody, response_json + ) + candidates: Final = payload.get("candidates", []) data_list: Final[list[ImageObject]] = [] for candidate in candidates: @@ -141,17 +154,18 @@ class GeminiImageEditConfig(BaseImageEditConfig): parts = content.get("parts", []) for part in parts: inline_data = part.get("inlineData") - if inline_data and inline_data.get("data"): + image_data = inline_data.get("data") if inline_data else None + if image_data: data_list.append( ImageObject( - b64_json=inline_data["data"], + b64_json=image_data, url=None, ) ) model_response.data = cast(list[OpenAIImage], data_list) - if "usageMetadata" in response_json: - model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"]) + if "usageMetadata" in payload: + model_response.usage = transform_gemini_image_usage(dict(decoded_json_object(payload["usageMetadata"]))) return model_response def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[dict[str, object]]: diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index fe82a9cae81..1f986b11af2 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -36,6 +37,7 @@ else: _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) _JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True) ) @@ -88,14 +90,12 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): """ complete_url: str = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL - complete_url = complete_url.rstrip("/") - # Gemini Flash Image Preview models use generateContent endpoint if is_gemini_image_model(model): - complete_url = f"{complete_url}/models/{model}:generateContent" + complete_url = join_api_base(complete_url, f"models/{model}:generateContent") else: # All other Imagen models use predict endpoint - complete_url = f"{complete_url}/models/{model}:predict" + complete_url = join_api_base(complete_url, f"models/{model}:predict") return complete_url @@ -189,7 +189,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): Transform Google AI Imagen response to litellm ImageResponse format """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", @@ -227,9 +227,9 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): ) # Extract usage metadata for Gemini models - if "usageMetadata" in response_data: + if "usageMetadata" in response_object: model_response.usage = transform_gemini_image_usage( - _JSON_DICT.validate_python(response_data["usageMetadata"]) + _JSON_DICT.validate_python(response_object["usageMetadata"]) ) web_search_requests: Final = get_gemini_image_web_search_requests(response_object) if web_search_requests and model_response.usage is not None: diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index a58037f691f..6f3e9a56b92 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -16,7 +16,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo from litellm.types.interactions import ( @@ -153,10 +153,8 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): if not api_key: raise ValueError("Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable.") - if stream: - return f"{api_base}/{self.api_version}/interactions?alt=sse" - - return f"{api_base}/{self.api_version}/interactions" + collection_url: Final = join_api_base(api_base, f"{self.api_version}/interactions") + return f"{collection_url}?alt=sse" if stream else collection_url def transform_request( self, @@ -304,7 +302,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): raise ValueError("Google API key is required") encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", + join_api_base(resolved_api_base, f"{self.api_version}/interactions/{encoded_interaction_id}"), {}, ) @@ -344,7 +342,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): raise ValueError("Google API key is required") encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", + join_api_base(resolved_api_base, f"{self.api_version}/interactions/{encoded_interaction_id}"), {}, ) @@ -375,7 +373,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): raise ValueError("Google API key is required") encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}/cancel", + join_api_base(resolved_api_base, f"{self.api_version}/interactions/{encoded_interaction_id}/cancel"), {}, ) diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index f4e2941670a..6f778ab2761 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -5,14 +5,16 @@ This file contains the transformation logic for the Gemini realtime API. import json from collections import OrderedDict from collections.abc import Mapping, Sequence -from typing import Any, Final, cast +from typing import Final, cast from typing_extensions import ReadOnly, Required, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -166,16 +168,17 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return True @staticmethod - def _usage_detail_alias(details: Mapping[str, int | None] | None, defaults: dict[str, int]) -> dict[str, int]: - if not isinstance(details, dict): + def _usage_detail_alias(details: object, defaults: dict[str, int]) -> dict[str, object]: + detail_mapping: Final = as_str_mapping(details) + if detail_mapping is None or not isinstance(details, dict): return dict(defaults) return { **defaults, - **{key: value for key, value in details.items() if value is not None}, + **{key: value for key, value in detail_mapping.items() if value is not None}, } @staticmethod - def _add_pipecat_usage_detail_aliases(usage_dict: dict[str, Any]) -> dict[str, object]: + def _add_pipecat_usage_detail_aliases(usage_dict: dict[str, object]) -> dict[str, object]: usage_dict.setdefault( "input_token_details", GeminiRealtimeConfig._usage_detail_alias( @@ -212,7 +215,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # so the API key must remain as a query parameter here. This is an accepted # limitation; httpx is not used for WebSocket so MaskedHTTPStatusError # already covers the main leak vector. - return f"{api_base}/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent?key={api_key}" + live_url: Final = join_api_base( + api_base, "ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent" + ) + return f"{live_url}?key={api_key}" def map_model_turn_event(self, model_turn: HttpxContentType) -> OpenAIRealtimeEventTypes: """ diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 34bd9660c43..11e619eb4c7 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.gemini.common_utils import ( GeminiError, @@ -141,15 +142,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): Note: This returns the base URL WITHOUT the API key. The API key will be appended to specific endpoint URLs in the transform methods. """ - if api_base is None: - api_base = GeminiModelInfo.get_api_base() - - if api_base is None: - raise ValueError("GEMINI_API_BASE is not set") - - # Ensure we're using the v1beta version for File Search - api_version: Final = "v1beta" - return f"{api_base}/{api_version}" + return join_api_base(GeminiModelInfo.get_api_base(api_base), "v1beta") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> GeminiError: """Return Gemini-specific error class.""" @@ -187,7 +180,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): api_key: Final = litellm_params.get("api_key") or GeminiModelInfo.get_api_key() if not api_key: raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required") - url: Final = f"{api_base}/models/{model}:generateContent" + url: Final = join_api_base(api_base, f"models/{model}:generateContent") # Build file_search tool configuration (using snake_case as per Gemini docs) file_search_config: Final[dict[str, object]] = {"file_search_store_names": [vector_store_id]} @@ -328,7 +321,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): """ Transform create request to Gemini's fileSearchStores format. """ - url: Final = f"{api_base}/fileSearchStores" + url: Final = join_api_base(api_base, "fileSearchStores") # API key is passed via x-goog-api-key header (set in validate_environment) diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 53ca5ed4adc..3f6b900fd01 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -4,10 +4,12 @@ from typing import TYPE_CHECKING, Any, Final import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.vertex_ai.videos.transformation import veo_video_count_from_parameters from litellm.secret_managers.main import get_secret_str @@ -36,6 +38,8 @@ else: LiteLLMLoggingObj = Any BaseLLMException = Any +_DECODED_JSON: Final = TypeAdapter(object) + def _convert_image_to_gemini_format(image_file) -> dict[str, str]: """ @@ -59,7 +63,7 @@ def _convert_image_to_gemini_format(image_file) -> dict[str, str]: def _json_payload(raw_response: httpx.Response) -> object: """Read an HTTP response body as an opaque JSON payload.""" - return raw_response.json() + return _DECODED_JSON.validate_python(raw_response.json()) def _usage_video_resolution_from_parameters( @@ -241,7 +245,7 @@ class GeminiVideoConfig(BaseVideoConfig): return api_base.rstrip("/") model_name: Final = model.replace("gemini/", "") - url: Final = f"{api_base.rstrip('/')}/v1beta/models/{model_name}:predictLongRunning" + url: Final = join_api_base(api_base, f"v1beta/models/{model_name}:predictLongRunning") return url @@ -377,7 +381,7 @@ class GeminiVideoConfig(BaseVideoConfig): GET https://generativelanguage.googleapis.com/v1beta/{operation_name} """ operation_name: Final = extract_original_video_id(video_id) - url: Final = f"{api_base.rstrip('/')}/v1beta/{operation_name}" + url: Final = join_api_base(api_base, f"v1beta/{operation_name}") params: Final[dict[str, object]] = {} return url, params @@ -451,7 +455,7 @@ class GeminiVideoConfig(BaseVideoConfig): """ operation_name: Final = extract_original_video_id(video_id) - status_url: Final = f"{api_base.rstrip('/')}/v1beta/{operation_name}" + status_url: Final = join_api_base(api_base, f"v1beta/{operation_name}") client: Final = litellm.module_level_client status_response: Final = client.get(url=status_url, headers=headers) status_response.raise_for_status() diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 1fe7de0ba85..65921b628fb 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -17,6 +17,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.gigachat.utils import convert_usage, get_api_base from litellm.secret_managers.main import get_secret_str @@ -108,8 +109,7 @@ class GigaChatConfig(BaseConfig): stream: bool | None = None, ) -> str: """Get complete API URL for chat completions.""" - base: Final = get_api_base(api_base) - return f"{base}/chat/completions" + return join_api_base(get_api_base(api_base), "chat/completions") def validate_environment( self, @@ -273,7 +273,7 @@ class GigaChatConfig(BaseConfig): # Default to None (don't set function_call) return None - def _upload_image(self, image_url: str) -> str | None: + def _upload_image(self, image_url: str) -> object: """ Upload image to GigaChat and return file_id. @@ -293,7 +293,7 @@ class GigaChatConfig(BaseConfig): verbose_logger.error("Failed to upload image: %s", e) return None - def _transform_list_content(self, content: Sequence) -> tuple[str, Sequence[str]]: + def _transform_list_content(self, content: Sequence) -> tuple[str, Sequence[object]]: """ Extract text and image attachments from a multimodal message content list. @@ -304,7 +304,7 @@ class GigaChatConfig(BaseConfig): Tuple of (combined text, list of attachment file ids) """ texts: Final[list[str]] = [] # mutable-ok: accumulator - attachments: Final[list[str]] = [] # mutable-ok: accumulator + attachments: Final[list[object]] = [] # mutable-ok: accumulator for part in content: if isinstance(part, dict): if part.get("type") == "text": diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index e0e746427ad..30ca7147865 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -16,6 +16,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.gigachat.utils import get_api_base @@ -98,8 +99,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): stream: bool | None = None, ) -> str: """Get the complete URL for embeddings endpoint.""" - base: Final = get_api_base(api_base) - return f"{base}/embeddings" + return join_api_base(get_api_base(api_base), "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index c051d693720..d53ae3af334 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -12,18 +12,23 @@ import uuid from collections.abc import Mapping from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, ) from litellm.llms.gigachat.utils import get_api_base +from litellm.llms.openai.common_utils import decoded_json_object from litellm.types.utils import LlmProviders from .authenticator import get_access_token, get_access_token_async # Simple in-memory cache for file IDs -_file_cache: Final[dict[str, str]] = {} +_file_cache: Final[dict[str, object]] = {} +_DECODED_JSON: Final = TypeAdapter(object) def _get_url_hash(url: str) -> str: @@ -90,7 +95,7 @@ def upload_file_sync( credentials: str | None = None, api_base: str | None = None, litellm_params: Mapping[str, object] | None = None, -) -> str | None: +) -> object: """ Upload file to GigaChat and return file_id (sync). @@ -125,8 +130,7 @@ def upload_file_sync( access_token: Final = get_access_token(credentials=credentials, litellm_params=litellm_params) # Upload to GigaChat - base_url: Final = get_api_base(api_base) - upload_url: Final = f"{base_url}/files" + upload_url: Final = join_api_base(get_api_base(api_base), "files") client: Final = get_httpx_client(params={"ssl_verify": False}) response: Final = client.post( @@ -137,9 +141,9 @@ def upload_file_sync( timeout=60, ) response.raise_for_status() - result: Final = response.json() + result: Final = _DECODED_JSON.validate_python(response.json()) - file_id: Final = result.get("id") + file_id: Final = decoded_json_object(result).get("id") if file_id: _file_cache[url_hash] = file_id verbose_logger.debug("File uploaded successfully, file_id: %s", file_id) @@ -156,7 +160,7 @@ async def upload_file_async( credentials: str | None = None, api_base: str | None = None, litellm_params: Mapping[str, object] | None = None, -) -> str | None: +) -> object: """ Upload file to GigaChat and return file_id (async). @@ -191,8 +195,7 @@ async def upload_file_async( access_token: Final = await get_access_token_async(credentials=credentials, litellm_params=litellm_params) # Upload to GigaChat - base_url: Final = get_api_base(api_base) - upload_url: Final = f"{base_url}/files" + upload_url: Final = join_api_base(get_api_base(api_base), "files") client: Final = get_async_httpx_client( llm_provider=LlmProviders.GIGACHAT, @@ -206,9 +209,9 @@ async def upload_file_async( timeout=60, ) response.raise_for_status() - result: Final = response.json() + result: Final = _DECODED_JSON.validate_python(response.json()) - file_id: Final = result.get("id") + file_id: Final = decoded_json_object(result).get("id") if file_id: _file_cache[url_hash] = file_id verbose_logger.debug("File uploaded successfully, file_id: %s", file_id) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index f16d263fd4d..27f75eece50 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -5,11 +5,13 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final import httpx +from pydantic import TypeAdapter from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.gigachat.authenticator import get_access_token from litellm.llms.gigachat.chat.streaming import GigaChatModelResponseIterator from litellm.llms.gigachat.utils import GIGACHAT_BASE_URL +from litellm.llms.openai.common_utils import decoded_json_object from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import EmbeddingResponse @@ -20,6 +22,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import CostResponseTypes +_DECODED_JSON: Final = TypeAdapter(object) + class GigaChatPassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: @@ -40,7 +44,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if base_target_url is None: raise Exception("GigaChat api base not found") - complete_url: Final = f"{base_target_url}/{endpoint.lstrip('/')}" + complete_url: Final = f"{base_target_url.rstrip('/')}/{endpoint.lstrip('/')}" return ( httpx.URL(complete_url), @@ -161,7 +165,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): continue chunk = chunk.removeprefix("data: ") try: - message = json.loads(chunk) + message = _DECODED_JSON.validate_python(json.loads(chunk)) except json.JSONDecodeError: continue @@ -169,7 +173,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): streaming_response=None, sync_stream=False, ) - translated_chunk = gigachat_iterator.chunk_parser(chunk=message) + translated_chunk = gigachat_iterator.chunk_parser(chunk=decoded_json_object(message)) if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields( # pyright: ignore[reportUnnecessaryIsInstance] # runtime guard for patched chunk_parser dict(translated_chunk) @@ -193,8 +197,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): @staticmethod def get_api_base(api_base: str | None = None) -> str | None: - resolved_api_base: Final = api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL - return resolved_api_base.rstrip("/") if resolved_api_base else None + return api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL @staticmethod def get_api_key( diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index ab1890ad73b..55ca4260221 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -21,5 +21,5 @@ def convert_usage(usage_data: Mapping[str, int]) -> Usage: ) -def get_api_base(api_base: str | None = None) -> str | None: +def get_api_base(api_base: str | None = None) -> str: return (api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL).rstrip("/") diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 9553327db97..83bc8df23a6 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -4,9 +4,11 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # spread dict loses the role literal without a cast import httpx +from pydantic import TypeAdapter from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.openai.common_utils import decoded_json_object from litellm.llms.openai.openai import OpenAIConfig from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.utils import ModelResponse @@ -23,6 +25,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.tokenizer import Encoding +_DECODED_JSON: Final = TypeAdapter(object) + class GithubCopilotConfig(OpenAIConfig): def __init__( @@ -309,7 +313,7 @@ class GithubCopilotConfig(OpenAIConfig): json_mode: bool | None = None, ) -> "ModelResponse": try: - response_json = raw_response.json() + decoded_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: return super().transform_response( model=model, @@ -325,12 +329,13 @@ class GithubCopilotConfig(OpenAIConfig): json_mode=json_mode, ) + response_json: Final = decoded_json_object(decoded_response) if not response_json.get("choices"): - response_json = self._synthesize_choices_for_anthropic_native(response_json) + synthesized: Final = self._synthesize_choices_for_anthropic_native(dict(response_json)) raw_response = httpx.Response( status_code=raw_response.status_code, headers=raw_response.headers, - content=json.dumps(response_json).encode(), + content=json.dumps(synthesized).encode(), ) return super().transform_response( diff --git a/litellm/llms/github_copilot/embedding/transformation.py b/litellm/llms/github_copilot/embedding/transformation.py index 7f9209afd65..4f308b01159 100644 --- a/litellm/llms/github_copilot/embedding/transformation.py +++ b/litellm/llms/github_copilot/embedding/transformation.py @@ -20,6 +20,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues @@ -121,7 +122,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): user_session: Final = require_github_copilot_user_session( cast("dict[str, object]", litellm_params) # cast-ok: litellm_params arrives as an untyped request dict ) - effective_api_base = ( + effective_api_base: Final = ( (user_session.api_base if user_session is not None else None) or api_base or self.authenticator.get_api_base() @@ -129,11 +130,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): or DEFAULT_GITHUB_COPILOT_API_BASE ) - # Remove trailing slashes - effective_api_base = effective_api_base.rstrip("/") - - # Return the embeddings endpoint - return f"{effective_api_base}/embeddings" + return join_api_base(effective_api_base, "embeddings") def transform_embedding_request( self, @@ -205,7 +202,9 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class(self, error_message: str, status_code: int, headers: Any) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers + ) -> BaseLLMException: from litellm.llms.openai.openai import OpenAIConfig return OpenAIConfig().get_error_class(error_message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/github_copilot/messages/transformation.py b/litellm/llms/github_copilot/messages/transformation.py index be138755f85..613b7e42100 100644 --- a/litellm/llms/github_copilot/messages/transformation.py +++ b/litellm/llms/github_copilot/messages/transformation.py @@ -1,6 +1,8 @@ -from typing import Any, Final, cast # noqa: TID251 # narrows untyped request dicts at the session boundary +from collections.abc import Sequence +from typing import Final, cast # noqa: TID251 # narrows untyped request dicts at the session boundary from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -53,7 +55,7 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig): self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, @@ -140,5 +142,5 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig): or api_base or self.authenticator.get_api_base() or DEFAULT_GITHUB_COPILOT_API_BASE - ).rstrip("/") - return resolved if resolved.endswith("/v1/messages") else f"{resolved}/v1/messages" + ) + return ensure_api_route(resolved, "v1/messages") diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index d468889b626..b055d388dd5 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -9,6 +9,7 @@ https://github.com/caozhiyuan/copilot-api """ import os +from collections.abc import Mapping from typing import ( TYPE_CHECKING, Any, @@ -16,10 +17,13 @@ from typing import ( cast, # noqa: TID251 # narrows untyped request dicts at the session boundary ) +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import ( ResponseInputParam, @@ -46,6 +50,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_OBJECT_LIST: Final = TypeAdapter(list[object]) +_OBJECT_MAPPING: Final = TypeAdapter(Mapping[object, object]) + def github_copilot_supports_responses_api(model: str) -> bool: """ @@ -266,7 +273,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): user_session: Final = require_github_copilot_user_session( cast("dict[str, object]", litellm_params) # cast-ok: litellm_params arrives as an untyped request dict ) - effective_api_base = ( + effective_api_base: Final = ( (user_session.api_base if user_session is not None else None) or api_base or self.authenticator.get_api_base() @@ -274,11 +281,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): or DEFAULT_GITHUB_COPILOT_API_BASE ) - # Remove trailing slashes - effective_api_base = effective_api_base.rstrip("/") - - # Return the responses endpoint - return f"{effective_api_base}/responses" + return join_api_base(effective_api_base, "responses") def _handle_reasoning_item(self, item: dict[str, object]) -> dict[str, object]: """ @@ -390,7 +393,9 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): """ return self._contains_vision_content(input_param) - def _contains_vision_content(self, value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bool: + def _contains_vision_content( + self, value: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH + ) -> bool: """ Recursively check if a value contains vision content. @@ -407,22 +412,28 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): # Check arrays if isinstance(value, list): - return any(self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in value) + items: Final = _OBJECT_LIST.validate_python(value) + return any(self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in items) # Only check dict/object types if not isinstance(value, dict): return False + fields: Final = _OBJECT_MAPPING.validate_python(value) + # Check if this item is an input_image - item_type: Final = value.get("type") + item_type: Final = fields.get("type") if isinstance(item_type, str) and item_type.lower() == "input_image": return True # Check content field recursively - if "content" in value and isinstance(value["content"], list): - return any( - self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in value["content"] - ) + if "content" in fields: + content: Final = fields["content"] + if isinstance(content, list): + content_items: Final = _OBJECT_LIST.validate_python(content) + return any( + self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in content_items + ) return False diff --git a/litellm/llms/gradient_ai/chat/transformation.py b/litellm/llms/gradient_ai/chat/transformation.py index fcee55a27e9..2a2368ad8de 100644 --- a/litellm/llms/gradient_ai/chat/transformation.py +++ b/litellm/llms/gradient_ai/chat/transformation.py @@ -1,5 +1,8 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final, Literal +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -42,7 +45,7 @@ class GradientAIConfig(OpenAILikeChatConfig): include_guardrails_info: bool | None = None, provide_citations: bool | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @@ -105,10 +108,10 @@ class GradientAIConfig(OpenAILikeChatConfig): ) -> str: clean_api_base: Final = (api_base or "").rstrip("/") if clean_api_base and clean_api_base != GRADIENT_AI_SERVERLESS_ENDPOINT: - return f"{clean_api_base}/api/v1/chat/completions" + return join_api_base(clean_api_base, "api/v1/chat/completions") agent_endpoint: Final = (get_secret_str("GRADIENT_AI_AGENT_ENDPOINT") or "").rstrip("/") if agent_endpoint and agent_endpoint != GRADIENT_AI_SERVERLESS_ENDPOINT: - return f"{agent_endpoint}/api/v1/chat/completions" + return join_api_base(agent_endpoint, "api/v1/chat/completions") return f"{GRADIENT_AI_SERVERLESS_ENDPOINT}/v1/chat/completions" def _get_openai_compatible_provider_info( diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index 8c269cb5dfd..23f68335ea9 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -11,7 +11,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields +from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_internal_message_fields from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, ) @@ -40,6 +40,7 @@ class GroqExecutedToolIdentity(LiteLLMBaseModel): _EXECUTED_TOOLS_ADAPTER: Final = TypeAdapter(tuple[GroqExecutedToolIdentity, ...]) +_ATTRIBUTE_VALUE: Final = TypeAdapter(object) class GroqChatConfig(OpenAILikeChatConfig): @@ -153,7 +154,7 @@ class GroqChatConfig(OpenAILikeChatConfig): for k, v in _message.items(): if v is not None: new_message[k] = v - messages[idx] = strip_litellm_response_only_fields(new_message) + messages[idx] = strip_litellm_internal_message_fields(new_message) if is_async: return super()._transform_messages(messages=messages, model=model, is_async=True) @@ -324,7 +325,7 @@ class GroqChatConfig(OpenAILikeChatConfig): json_mode=json_mode, ) - original_service_tier: Final[object] = getattr(model_response, "service_tier") + original_service_tier: Final = _ATTRIBUTE_VALUE.validate_python(getattr(model_response, "service_tier")) mapped_service_tier: Final[Literal["auto", "default", "flex"]] = self._map_groq_service_tier( original_service_tier=original_service_tier ) diff --git a/litellm/llms/groq/cost_calculator.py b/litellm/llms/groq/cost_calculator.py index bcebb347aea..7f75a607410 100644 --- a/litellm/llms/groq/cost_calculator.py +++ b/litellm/llms/groq/cost_calculator.py @@ -19,7 +19,10 @@ if TYPE_CHECKING: def cost_per_web_search_request(usage: Usage, model_info: "ModelInfo") -> float: search_costs: Final = model_info.get("search_context_cost_per_query") cost_per_search: Final = search_costs.get("search_context_size_medium", 0.0) if search_costs else 0.0 - server_tool_use: Final = getattr(usage, "server_tool_use", None) + try: + server_tool_use: Final = usage.server_tool_use + except AttributeError: + return 0.0 if server_tool_use is None: return 0.0 searches: Final = server_tool_use.web_search_requests or 0 diff --git a/litellm/llms/groq/stt/transformation.py b/litellm/llms/groq/stt/transformation.py index 93d36dd7b0f..33e01ab7c0e 100644 --- a/litellm/llms/groq/stt/transformation.py +++ b/litellm/llms/groq/stt/transformation.py @@ -3,6 +3,7 @@ Translate from OpenAI's `/v1/audio/transcriptions` to Groq's `/v1/audio/transcri """ import types +from collections.abc import Mapping from typing import Final import litellm @@ -39,16 +40,17 @@ class GroqSTTConfig: tools: list | None = None, tool_choice: str | dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/heroku/chat/transformation.py b/litellm/llms/heroku/chat/transformation.py index acd630c85d0..002078dd04b 100644 --- a/litellm/llms/heroku/chat/transformation.py +++ b/litellm/llms/heroku/chat/transformation.py @@ -6,11 +6,12 @@ this is OpenAI compatible - no translation needed / occurs import os from collections.abc import Coroutine -from typing import Final, Literal, overload +from typing import Literal, overload from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import AllMessageValues @@ -78,9 +79,4 @@ class HerokuChatConfig(OpenAIGPTConfig): "No api base was set. Please provide an api_base, or set the HEROKU_API_BASE environment variable." ) - clean_api_base: Final = resolved_base.rstrip("/") - - if not clean_api_base.endswith("/v1/chat/completions"): - return f"{clean_api_base}/v1/chat/completions" - - return clean_api_base + return ensure_api_route(resolved_base, "v1/chat/completions") diff --git a/litellm/llms/hosted_vllm/embedding/transformation.py b/litellm/llms/hosted_vllm/embedding/transformation.py index 3b54145a67b..d72a8d06f21 100644 --- a/litellm/llms/hosted_vllm/embedding/transformation.py +++ b/litellm/llms/hosted_vllm/embedding/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -81,14 +82,7 @@ class HostedVLLMEmbeddingConfig(BaseEmbeddingConfig): if api_base is None: raise ValueError("api_base is required for hosted_vllm embeddings") - # Remove trailing slashes - api_base = api_base.rstrip("/") - - # Ensure the URL ends with /embeddings - if not api_base.endswith("/embeddings"): - api_base = f"{api_base}/embeddings" - - return api_base + return ensure_api_route(api_base, "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py index 6804d3a0fb8..8e05b6dfb46 100644 --- a/litellm/llms/hosted_vllm/image_edit/transformation.py +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -1,5 +1,6 @@ from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -39,5 +40,5 @@ class HostedVLLMImageEditConfig(OpenAIImageEditConfig): ) trimmed: Final = resolved_api_base.rstrip("/") if trimmed.endswith("/v1"): - return f"{trimmed}/images/edits" - return f"{trimmed}/v1/images/edits" + return join_api_base(trimmed, "images/edits") + return join_api_base(trimmed, "v1/images/edits") diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index a463fc1d3c9..0b77a04c69a 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -7,13 +7,15 @@ from types import MappingProxyType from typing import Final import httpx -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError from litellm._uuid import uuid from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.llms.openai.common_utils import decoded_json_object from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( HostedVLLMRerankTruncationParams, @@ -27,6 +29,8 @@ from litellm.types.rerank import ( RerankTokens, ) +_DECODED_JSON: Final = TypeAdapter(object) + class HostedVLLMRerankError(BaseLLMException): def __init__( @@ -62,7 +66,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if api_base.endswith("/v1/rerank"): api_base = api_base.replace("/v1/rerank", "/rerank") elif not api_base.endswith("/rerank"): - api_base = f"{api_base}/rerank" + api_base = join_api_base(api_base, "rerank") return api_base raise ValueError("api_base must be provided for Hosted VLLM rerank") @@ -192,11 +196,11 @@ class HostedVLLMRerankConfig(BaseRerankConfig): Process response from Hosted VLLM rerank API """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raise ValueError(f"Error parsing response: {raw_response.text}, status_code={raw_response.status_code}") - return self._transform_response(raw_response_json) + return self._transform_response(dict(decoded_json_object(raw_response_json))) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return HostedVLLMRerankError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/hosted_vllm/responses/transformation.py b/litellm/llms/hosted_vllm/responses/transformation.py index 20d755a2273..8489bc0da84 100644 --- a/litellm/llms/hosted_vllm/responses/transformation.py +++ b/litellm/llms/hosted_vllm/responses/transformation.py @@ -8,6 +8,7 @@ the chat completions → responses conversion pipeline. from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -64,9 +65,9 @@ class HostedVLLMResponsesAPIConfig(OpenAIResponsesAPIConfig): # If api_base already ends with /v1, append /responses # Otherwise append /v1/responses if api_base.endswith("/v1"): - return f"{api_base}/responses" + return join_api_base(api_base, "responses") - return f"{api_base}/v1/responses" + return join_api_base(api_base, "v1/responses") def supports_native_websocket(self) -> bool: """Hosted vLLM does not support native WebSocket for Responses API""" diff --git a/litellm/llms/hosted_vllm/transcriptions/transformation.py b/litellm/llms/hosted_vllm/transcriptions/transformation.py index 9a8c67fadc0..785ccd52a69 100644 --- a/litellm/llms/hosted_vllm/transcriptions/transformation.py +++ b/litellm/llms/hosted_vllm/transcriptions/transformation.py @@ -6,6 +6,7 @@ from typing import Final import httpx +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, ) @@ -40,11 +41,7 @@ class HostedVLLMAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig): stream: bool | None = None, ) -> str: if api_base: - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/audio/transcriptions"): - api_base = f"{api_base}/v1/audio/transcriptions" - return api_base + return ensure_api_route(api_base, "v1/audio/transcriptions") raise ValueError("api_base must be provided for Hosted VLLM rerank") def transform_audio_transcription_request( diff --git a/litellm/llms/hosted_vllm/videos/transformation.py b/litellm/llms/hosted_vllm/videos/transformation.py index 22e4f4876fb..2045d2c8988 100644 --- a/litellm/llms/hosted_vllm/videos/transformation.py +++ b/litellm/llms/hosted_vllm/videos/transformation.py @@ -8,14 +8,16 @@ from typing import Final from urllib.parse import urlparse from httpx._types import FileTypes, RequestFiles +from pydantic import TypeAdapter from litellm.images.utils import ImageEditRequestUtils -from litellm.litellm_core_utils.url_utils import SSRFError, validate_url +from litellm.litellm_core_utils.url_utils import SSRFError, join_api_base, validate_url from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import VideoCreateOptionalRequestParams +_DECODED_JSON: Final = TypeAdapter(object) _EXCLUDED_FORM_KEYS: Final = frozenset( { "model", @@ -82,7 +84,7 @@ def _maybe_json(value: object) -> object: stripped: Final = value.strip() if not stripped or stripped[0] not in "{[": return value - return json.loads(stripped) + return _DECODED_JSON.validate_python(json.loads(stripped)) def _reject_unsafe_media_url(url: str) -> None: @@ -177,8 +179,8 @@ class HostedVLLMVideoConfig(OpenAIVideoConfig): ) trimmed: Final = resolved_api_base.rstrip("/") if trimmed.endswith("/v1"): - return f"{trimmed}/videos" - return f"{trimmed}/v1/videos" + return join_api_base(trimmed, "videos") + return join_api_base(trimmed, "v1/videos") def transform_video_create_request( self, diff --git a/litellm/llms/huggingface/chat/transformation.py b/litellm/llms/huggingface/chat/transformation.py index b45c706f8bf..3fcb738eacb 100644 --- a/litellm/llms/huggingface/chat/transformation.py +++ b/litellm/llms/huggingface/chat/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest if TYPE_CHECKING: @@ -29,11 +30,11 @@ def _build_chat_completion_url(model_url: str) -> str: # Append /chat/completions if not already present if model_url.endswith("/v1"): - model_url += "/chat/completions" + model_url = join_api_base(model_url, "chat/completions") # Append /v1/chat/completions if not already present if not model_url.endswith("/chat/completions"): - model_url += "/v1/chat/completions" + model_url = join_api_base(model_url, "v1/chat/completions") return model_url diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 60917c68221..a1d1a3f6a19 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -3,6 +3,7 @@ import os from typing import Final, Literal, Protocol, get_args import httpx +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -11,6 +12,7 @@ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, ) +from litellm.llms.openai.common_utils import decoded_json_object from litellm.types.utils import EmbeddingResponse from ...base import BaseLLM @@ -21,6 +23,8 @@ config: Final = HuggingFaceEmbeddingConfig() HF_HUB_URL: Final = "https://huggingface.co" +_DECODED_JSON: Final = TypeAdapter(object) + hf_tasks_embeddings: Final = ( Literal[ # pipeline tags + hf tei endpoints - https://huggingface.github.io/text-embeddings-inference/#/ "sentence-similarity", "feature-extraction", "rerank", "embed", "similarity" @@ -34,7 +38,7 @@ class _SupportsTokenEncode(Protocol): def encode(self, text: str) -> list[int]: ... -def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> str | None: +def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type @@ -44,14 +48,14 @@ def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: model_info: Final = http_client.get(url=f"{api_base}/api/models/{model}") - model_info_dict: Final = model_info.json() + model_info_dict: Final = decoded_json_object(_DECODED_JSON.validate_python(model_info.json())) - pipeline_tag: Final[str | None] = model_info_dict.get("pipeline_tag", None) + pipeline_tag: Final = model_info_dict.get("pipeline_tag", None) return pipeline_tag -async def async_get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> str | None: +async def async_get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type @@ -63,9 +67,9 @@ async def async_get_hf_task_embedding_for_model(model: str, task_type: str | Non model_info: Final = await http_client.get(url=f"{api_base}/api/models/{model}") - model_info_dict: Final = model_info.json() + model_info_dict: Final = decoded_json_object(_DECODED_JSON.validate_python(model_info.json())) - pipeline_tag: Final[str | None] = model_info_dict.get("pipeline_tag", None) + pipeline_tag: Final = model_info_dict.get("pipeline_tag", None) return pipeline_tag @@ -77,7 +81,7 @@ class HuggingFaceEmbedding(BaseLLM): def __init__(self) -> None: super().__init__() - def _transform_input_on_pipeline_tag(self, input: list, pipeline_tag: str | None) -> dict: + def _transform_input_on_pipeline_tag(self, input: list, pipeline_tag: object) -> dict: if pipeline_tag is None: return {"inputs": input} if pipeline_tag == "sentence-similarity" or pipeline_tag == "similarity": diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 0883a9a8077..1b7d6483739 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -447,14 +447,13 @@ class HuggingFaceEmbeddingConfig(BaseConfig): except Exception: # this should remain non blocking we should not block a response returning if calculating usage fails pass - output_text: Final = model_response["choices"][0]["message"].get("content", "") + first_message: Final = model_response.choices[0].message + output_text: Final = first_message.content if output_text is not None and len(output_text) > 0: completion_tokens = 0 try: if encoding is not None: - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) - ) ##[TODO] use the llama2 tokenizer here + completion_tokens = len(encoding.encode(output_text)) ##[TODO] use the llama2 tokenizer here except Exception: # this should remain non blocking we should not block a response returning if calculating usage fails pass diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 334c60ee848..1f130f84283 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -7,6 +7,7 @@ from typing_extensions import TypedDict import litellm from litellm._uuid import uuid +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -69,15 +70,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): """ Get the complete URL for the API call, including the /rerank suffix if necessary. """ - # Get base URL from api_base or default - base_url = self.get_api_base(model=model, api_base=api_base) - - # Remove trailing slashes and ensure we have the /rerank endpoint - base_url = base_url.rstrip("/") - if not base_url.endswith("/rerank"): - base_url = f"{base_url}/rerank" - - return base_url + return ensure_api_route(self.get_api_base(model=model, api_base=api_base), "rerank") def get_supported_cohere_rerank_params(self, model: str) -> list: return [ @@ -93,7 +86,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/llms/infinity/embedding/transformation.py b/litellm/llms/infinity/embedding/transformation.py index f289918cb3c..a467c56fc74 100644 --- a/litellm/llms/infinity/embedding/transformation.py +++ b/litellm/llms/infinity/embedding/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -36,11 +37,7 @@ class InfinityEmbeddingConfig(BaseEmbeddingConfig): ) -> str: if api_base is None: raise ValueError("api_base is required for Infinity embeddings") - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/embeddings"): - api_base = f"{api_base}/embeddings" - return api_base + return ensure_api_route(api_base, "embeddings") def validate_environment( self, diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index e089b3fecbe..07f3bb2c26b 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -13,6 +13,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( @@ -61,11 +62,7 @@ class InfinityRerankConfig(CohereRerankConfig): ) -> str: if api_base is None: raise ValueError("api_base is required for Infinity rerank") - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/rerank"): - api_base = f"{api_base}/rerank" - return api_base + return ensure_api_route(api_base, "rerank") def validate_environment( self, diff --git a/litellm/llms/jina_ai/embedding/transformation.py b/litellm/llms/jina_ai/embedding/transformation.py index 9b156e5ae21..04ff8f941a3 100644 --- a/litellm/llms/jina_ai/embedding/transformation.py +++ b/litellm/llms/jina_ai/embedding/transformation.py @@ -15,6 +15,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm import BaseEmbeddingConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -115,12 +116,7 @@ class JinaAIEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - if api_base: - clean_api_base: Final = api_base.rstrip("/") - if not clean_api_base.endswith("/embeddings"): - return f"{clean_api_base}/embeddings" - return clean_api_base - return "https://api.jina.ai/v1/embeddings" + return ensure_api_route(api_base or "https://api.jina.ai/v1", "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/langflow/chat/transformation.py b/litellm/llms/langflow/chat/transformation.py index 1781082d36f..3807be562cf 100644 --- a/litellm/llms/langflow/chat/transformation.py +++ b/litellm/llms/langflow/chat/transformation.py @@ -1,15 +1,19 @@ """LangFlow run API: POST {api_base}/api/v1/run/{flow_id}""" +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from urllib.parse import quote import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.openai.common_utils import decoded_json_object from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -26,6 +30,8 @@ else: AsyncHTTPHandler = Any CustomStreamWrapper = Any +_DECODED_JSON: Final = TypeAdapter(object) + class LangFlowError(BaseLLMException): """Exception class for LangFlow API errors.""" @@ -107,9 +113,8 @@ class LangFlowConfig(BaseConfig): "api_base is required for LangFlow. Set it via LANGFLOW_API_BASE env var or api_base parameter." ) - api_base = api_base.rstrip("/") flow_id: Final = quote(self._get_flow_id(model, optional_params), safe="") - return f"{api_base}/api/v1/run/{flow_id}" + return join_api_base(api_base, f"api/v1/run/{flow_id}") def _get_last_user_message(self, messages: list[AllMessageValues]) -> str: """Extract the text of the last user message to use as input_value.""" @@ -179,7 +184,7 @@ class LangFlowConfig(BaseConfig): verbose_logger.debug("LangFlow request payload: %s", payload) return payload - def _extract_content_from_response(self, response_json: dict) -> str | None: + def _extract_content_from_response(self, response_json: Mapping[str, object]) -> str | None: """ Extract the assistant text from a LangFlow run response. @@ -236,7 +241,7 @@ class LangFlowConfig(BaseConfig): json_mode: bool | None = None, ) -> ModelResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise LangFlowError( message=f"LangFlow returned a non-JSON response: {e}", @@ -245,7 +250,7 @@ class LangFlowConfig(BaseConfig): verbose_logger.debug("LangFlow response: %s", response_json) - content: Final = self._extract_content_from_response(response_json) + content: Final = self._extract_content_from_response(decoded_json_object(response_json)) if content is None: raise LangFlowError( message=( diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 499cbab3b7b..575820bceee 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -9,10 +9,13 @@ import uuid from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices +_DECODED_JSON: Final = TypeAdapter(object) + class LangGraphSSEStreamIterator: """ @@ -60,7 +63,7 @@ class LangGraphSSEStreamIterator: return None try: - data: Final = json.loads(json_str) + data: Final = _DECODED_JSON.validate_python(json.loads(json_str)) return self._process_data(data) except json.JSONDecodeError: verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100]) diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index 4f93f435957..603d5a07292 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -13,13 +13,16 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional, Union import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator +from litellm.llms.openai.common_utils import decoded_json_object from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -36,6 +39,8 @@ else: AsyncHTTPHandler = Any CustomStreamWrapper = Any +_DECODED_JSON: Final = TypeAdapter(object) + class LangGraphError(BaseLLMException): """Exception class for LangGraph API errors.""" @@ -116,14 +121,11 @@ class LangGraphConfig(BaseConfig): "api_base is required for LangGraph. Set it via LANGGRAPH_API_BASE env var or api_base parameter." ) - # Remove trailing slash if present - api_base = api_base.rstrip("/") - # Choose endpoint based on streaming mode if stream: - return f"{api_base}/runs/stream" + return join_api_base(api_base, "runs/stream") else: - return f"{api_base}/runs/wait" + return join_api_base(api_base, "runs/wait") def _get_assistant_id(self, model: str, optional_params: dict) -> str: """ @@ -235,7 +237,7 @@ class LangGraphConfig(BaseConfig): verbose_logger.debug("LangGraph request payload: %s", payload) return payload - def _extract_content_from_response(self, response_json: dict) -> str: + def _extract_content_from_response(self, response_json: Mapping[str, object]) -> str: """ Extract content from LangGraph non-streaming response. @@ -436,10 +438,10 @@ class LangGraphConfig(BaseConfig): Transform the LangGraph response to LiteLLM ModelResponse format. """ try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) verbose_logger.debug("LangGraph response: %s", response_json) - content: Final = self._extract_content_from_response(response_json) + content: Final = self._extract_content_from_response(decoded_json_object(response_json)) # Create the message message: Final = Message(content=content, role="assistant") diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index a740050dd94..42bd63b0160 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -3,7 +3,7 @@ Translate from OpenAI's `/v1/chat/completions` to Lemonade's `/v1/chat/completio """ from collections.abc import Iterable, Mapping -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import quote import httpx @@ -12,6 +12,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -26,6 +27,7 @@ if TYPE_CHECKING: _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _STR: Final = TypeAdapter(str) +_OPAQUE_VALUE: Final = TypeAdapter(object) class _JsonResponse(Protocol): @@ -105,7 +107,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): # Getting the list of models from lemonade try: response: Final = litellm.module_level_client.get( - url=f"{api_base}/models", + url=join_api_base(api_base, "models"), headers=self._get_auth_headers(api_key), ) except Exception as e: @@ -122,7 +124,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): return ["lemonade/" + _STR.validate_python(model["id"]) for model in model_list] @staticmethod - def _get_positive_int(value: Any) -> int | None: + def _get_positive_int(value: object) -> int | None: if isinstance(value, bool): return None if isinstance(value, int) and value > 0: @@ -162,7 +164,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): provider_specific_entry.get("context_window"), provider_specific_entry.get("max_context_window"), ): - parsed = self._get_positive_int(value) + parsed = self._get_positive_int(_OPAQUE_VALUE.validate_python(value)) if parsed is not None: return parsed return None @@ -193,7 +195,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): try: response: Final = litellm.module_level_client.get( - url=f"{api_base}/models/{encoded_model}", + url=join_api_base(str(api_base), f"models/{encoded_model}"), headers=self._get_auth_headers(api_key), ) response.raise_for_status() diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py index e219a227b1f..fd1ca792302 100644 --- a/litellm/llms/linkup/search/transformation.py +++ b/litellm/llms/linkup/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -87,11 +88,7 @@ class LinkupSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("LINKUP_API_BASE") or self.LINKUP_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("LINKUP_API_BASE") or self.LINKUP_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/litellm_proxy/image_edit/transformation.py b/litellm/llms/litellm_proxy/image_edit/transformation.py index 9e5a2a2c343..f319e31ff3f 100644 --- a/litellm/llms/litellm_proxy/image_edit/transformation.py +++ b/litellm/llms/litellm_proxy/image_edit/transformation.py @@ -1,3 +1,4 @@ +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -21,5 +22,4 @@ class LiteLLMProxyImageEditConfig(OpenAIImageEditConfig): api_base = api_base or get_secret_str("LITELLM_PROXY_API_BASE") if api_base is None: raise ValueError("api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`") - api_base = api_base.rstrip("/") - return f"{api_base}/images/edits" + return join_api_base(api_base, "images/edits") diff --git a/litellm/llms/litellm_proxy/image_generation/transformation.py b/litellm/llms/litellm_proxy/image_generation/transformation.py index fcddeabfc85..05e536eece5 100644 --- a/litellm/llms/litellm_proxy/image_generation/transformation.py +++ b/litellm/llms/litellm_proxy/image_generation/transformation.py @@ -1,3 +1,4 @@ +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.image_generation.gpt_transformation import ( GPTImageGenerationConfig, ) @@ -33,5 +34,4 @@ class LiteLLMProxyImageGenerationConfig(GPTImageGenerationConfig): api_base = api_base or get_secret_str("LITELLM_PROXY_API_BASE") if api_base is None: raise ValueError("api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`") - api_base = api_base.rstrip("/") - return f"{api_base}/images/generations" + return join_api_base(api_base, "images/generations") diff --git a/litellm/llms/litellm_proxy/responses/transformation.py b/litellm/llms/litellm_proxy/responses/transformation.py index 2928b07665a..3092b49b743 100644 --- a/litellm/llms/litellm_proxy/responses/transformation.py +++ b/litellm/llms/litellm_proxy/responses/transformation.py @@ -5,6 +5,7 @@ LiteLLM Proxy supports the OpenAI Responses API natively when the underlying mod This config enables pass-through behavior to the proxy's /v1/responses endpoint. """ +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.utils import LlmProviders @@ -40,10 +41,7 @@ class LiteLLMProxyResponsesAPIConfig(OpenAIResponsesAPIConfig): "Set via api_base parameter or LITELLM_PROXY_API_BASE environment variable" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def supports_native_websocket(self) -> bool: """LiteLLM Proxy does not support native WebSocket for Responses API""" diff --git a/litellm/llms/litellm_proxy/skills/prompt_injection.py b/litellm/llms/litellm_proxy/skills/prompt_injection.py index 3377338a03e..e8e1c68a112 100644 --- a/litellm/llms/litellm_proxy/skills/prompt_injection.py +++ b/litellm/llms/litellm_proxy/skills/prompt_injection.py @@ -9,9 +9,10 @@ import posixpath import zipfile from collections.abc import Mapping from io import BytesIO -from typing import Any, Final +from typing import Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.models.skills import LiteLLM_SkillsTable @@ -251,27 +252,30 @@ class SkillPromptInjectionHandler: if len(description) > max_desc_length: description = description[: max_desc_length - 3] + "..." - tool: Final[dict[str, Any]] = { + params: Final = ( + as_str_mapping(skill.metadata.get("parameters")) + if skill.metadata and isinstance(skill.metadata, dict) + else None + ) + parameters: Final[Mapping[str, object]] = ( + params + if params and isinstance(params, dict) + else { + "type": "object", + "properties": {}, + "required": [], + } + ) + + return { "type": "function", "function": { "name": func_name, "description": description, - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, + "parameters": parameters, }, } - # If skill has metadata with parameter definitions, use them - if skill.metadata and isinstance(skill.metadata, dict): - params: Final = skill.metadata.get("parameters") - if params and isinstance(params, dict): - tool["function"]["parameters"] = params - - return tool - def convert_skill_to_anthropic_tool(self, skill: LiteLLM_SkillsTable) -> Mapping[str, object]: """ Convert a LiteLLM skill to an Anthropic-style tool (messages API format). diff --git a/litellm/llms/litellm_proxy/skills/sandbox_executor.py b/litellm/llms/litellm_proxy/skills/sandbox_executor.py index 3e38dd81905..8e5c9effa17 100644 --- a/litellm/llms/litellm_proxy/skills/sandbox_executor.py +++ b/litellm/llms/litellm_proxy/skills/sandbox_executor.py @@ -7,7 +7,8 @@ Supports Docker, Podman, and Kubernetes backends. import base64 import os -from typing import Any, Final, Protocol, TypedDict +from collections.abc import Sequence +from typing import Final, Protocol, TypedDict from typing_extensions import ReadOnly @@ -41,6 +42,13 @@ class _GeneratedFile(TypedDict): mime_type: ReadOnly[str] +class _SandboxExecution(TypedDict): + success: ReadOnly[bool] + output: ReadOnly[str] + error: ReadOnly[str] + files: ReadOnly[Sequence[_GeneratedFile]] + + class SkillsSandboxExecutor: """ Executes skill code in llm-sandbox Docker container. @@ -76,7 +84,7 @@ class SkillsSandboxExecutor: code: str, skill_files: dict[str, bytes], requirements: str | None = None, - ) -> dict[str, Any]: + ) -> _SandboxExecution: """ Execute code with skill files in sandbox. diff --git a/litellm/llms/lm_studio/embed/transformation.py b/litellm/llms/lm_studio/embed/transformation.py index 66b2dbb534d..d6264ceb6e0 100644 --- a/litellm/llms/lm_studio/embed/transformation.py +++ b/litellm/llms/lm_studio/embed/transformation.py @@ -7,6 +7,7 @@ Docs - https://lmstudio.ai/docs/basics/server """ import types +from collections.abc import Mapping from typing import Final @@ -18,16 +19,17 @@ class LmStudioEmbeddingConfig: def __init__( self, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/manus/files/transformation.py b/litellm/llms/manus/files/transformation.py index 94b0e62199d..991aa4cbc32 100644 --- a/litellm/llms/manus/files/transformation.py +++ b/litellm/llms/manus/files/transformation.py @@ -21,7 +21,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, @@ -140,8 +140,8 @@ class ManusFilesConfig(BaseFilesConfig): # Manus API uses /v1/files endpoint if api_base.endswith("/v1"): - return f"{api_base}/files" - return f"{api_base}/v1/files" + return join_api_base(api_base, "files") + return join_api_base(api_base, "v1/files") def get_error_class( self, @@ -423,7 +423,7 @@ class ManusFilesConfig(BaseFilesConfig): litellm_params=litellm_params, ) encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{api_base}/{encoded_file_id}/content", {} + return join_api_base(api_base, f"{encoded_file_id}/content"), {} def transform_file_content_response( self, diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index 9785a50bbce..ed06391a260 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -10,7 +10,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( safe_convert_created_field, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str @@ -119,8 +119,8 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): # Manus API uses /v1/responses endpoint (OpenAI-compatible) if api_base.endswith("/v1"): - return f"{api_base}/responses" - return f"{api_base}/v1/responses" + return join_api_base(api_base, "responses") + return join_api_base(api_base, "v1/responses") def transform_responses_api_request( self, diff --git a/litellm/llms/maritalk.py b/litellm/llms/maritalk.py index 094dee4cee1..36bde5847fb 100644 --- a/litellm/llms/maritalk.py +++ b/litellm/llms/maritalk.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final from httpx._models import Headers @@ -32,7 +33,7 @@ class MaritalkConfig(OpenAIGPTConfig): tools: list[dict] | None = None, tool_choice: str | dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 70d3649debd..14470922e57 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -4,11 +4,14 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, ) +from litellm.llms.openai.common_utils import decoded_json_object from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -29,6 +32,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + MILVUS_OPTIONAL_PARAMS: Final = { "annsField", "limit", @@ -185,7 +190,7 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): } litellm_logging_obj.model_call_details["input"] = query_text litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") - return f"{api_base}/v2/vectordb/entities/search", { + return join_api_base(api_base, "v2/vectordb/entities/search"), { "collectionName": vector_store_id, "data": [query_vector], "annsField": "book_intro_vector", @@ -217,15 +222,17 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): results: Final = response_json.get("data", []) # Try to get text_field from optional_params first, then litellm_params - optional_params: Final = litellm_logging_obj.model_call_details.get("optional_params", {}) + optional_params: Final = decoded_json_object( + _DECODED_JSON.validate_python(litellm_logging_obj.model_call_details.get("optional_params", {})) + ) text_field = optional_params.get("milvus_text_field", "") # Fallback to litellm_params if not in optional_params if not text_field: - text_field = litellm_logging_obj.model_call_details.get("litellm_params", {}).get( - "milvus_text_field", "" - ) + text_field = decoded_json_object( + _DECODED_JSON.validate_python(litellm_logging_obj.model_call_details.get("litellm_params", {})) + ).get("milvus_text_field", "") # Transform results to standard format search_results: Final[list[VectorStoreSearchResult]] = [] diff --git a/litellm/llms/minimax/chat/transformation.py b/litellm/llms/minimax/chat/transformation.py index c5aa8811f02..f397bf7fbb9 100644 --- a/litellm/llms/minimax/chat/transformation.py +++ b/litellm/llms/minimax/chat/transformation.py @@ -5,6 +5,7 @@ MiniMax OpenAI transformation config - extends OpenAI chat config for MiniMax's from typing import Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -55,17 +56,14 @@ class MinimaxChatConfig(OpenAIGPTConfig): Override to ensure we use MiniMax's endpoint. """ # Get the base URL (either provided or default MiniMax endpoint) - base_url: Final = self.get_api_base(api_base=api_base) + base_url: Final = self.get_api_base(api_base=api_base).rstrip("/") # Ensure it ends with /chat/completions if base_url.endswith("/chat/completions"): return base_url - elif base_url.endswith("/v1"): - return f"{base_url}/chat/completions" - elif base_url.endswith("/"): - return f"{base_url}v1/chat/completions" - else: - return f"{base_url}/v1/chat/completions" + if base_url.endswith("/v1"): + return join_api_base(base_url, "chat/completions") + return join_api_base(base_url, "v1/chat/completions") def remove_cache_control_flag_from_messages_and_tools( self, diff --git a/litellm/llms/minimax/messages/transformation.py b/litellm/llms/minimax/messages/transformation.py index d62b88a24c6..f7a263a467c 100644 --- a/litellm/llms/minimax/messages/transformation.py +++ b/litellm/llms/minimax/messages/transformation.py @@ -2,9 +2,8 @@ MiniMax Anthropic transformation config - extends AnthropicConfig for MiniMax's Anthropic-compatible API """ -from typing import Any, Final # noqa: TID251 # override below must mirror the legacy base signature - import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -53,7 +52,7 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): self, headers: dict, # mutable-ok: mirrors the legacy base override signature model: str, - messages: list[Any], # mutable-ok: mirrors the legacy base override signature + messages: list[object], # mutable-ok: mirrors the legacy base override signature optional_params: dict, # mutable-ok: mirrors the legacy base override signature litellm_params: dict, # mutable-ok: mirrors the legacy base override signature api_key: str | None = None, @@ -82,15 +81,4 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): Get the complete URL for MiniMax API. Override to ensure we use MiniMax's endpoint, not Anthropic's. """ - # Get the base URL (either provided or default MiniMax endpoint) - base_url: Final = self.get_api_base(api_base=api_base) - - # If the base URL already includes the full path, return it - if base_url.endswith("/v1/messages"): - return base_url - - # Otherwise append the messages endpoint - if base_url.endswith("/"): - return f"{base_url}v1/messages" - else: - return f"{base_url}/v1/messages" + return ensure_api_route(self.get_api_base(api_base=api_base), "v1/messages") diff --git a/litellm/llms/minimax/text_to_speech/transformation.py b/litellm/llms/minimax/text_to_speech/transformation.py index 1cabc43eb1c..6f2cb36ed03 100644 --- a/litellm/llms/minimax/text_to_speech/transformation.py +++ b/litellm/llms/minimax/text_to_speech/transformation.py @@ -13,6 +13,7 @@ from httpx import Headers from pydantic import ConfigDict, TypeAdapter import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, @@ -404,10 +405,5 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): """ Construct the MiniMax endpoint URL. """ - base_url = api_base or get_secret_str("MINIMAX_API_BASE") or self.TTS_BASE_URL - base_url = base_url.rstrip("/") - - # MiniMax uses a simple endpoint path - url: Final = f"{base_url}{self.TTS_ENDPOINT_PATH}" - - return url + base_url: Final = api_base or get_secret_str("MINIMAX_API_BASE") or self.TTS_BASE_URL + return join_api_base(base_url, self.TTS_ENDPOINT_PATH) diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py index 5a8960d0898..ac0e1a33620 100644 --- a/litellm/llms/mistral/audio_speech/transformation.py +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Final import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, @@ -118,9 +119,7 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> str: - configured_base: Final = (api_base or self.TTS_BASE_URL).rstrip("/") - versioned_base: Final = configured_base if configured_base.endswith("/v1") else f"{configured_base}/v1" - return f"{versioned_base}/audio/speech" + return join_api_base(ensure_api_route(api_base or self.TTS_BASE_URL, "v1"), "audio/speech") def transform_text_to_speech_request( self, diff --git a/litellm/llms/mistral/audio_transcription/transformation.py b/litellm/llms/mistral/audio_transcription/transformation.py index 4fce7a95350..3eb9e122265 100644 --- a/litellm/llms/mistral/audio_transcription/transformation.py +++ b/litellm/llms/mistral/audio_transcription/transformation.py @@ -9,6 +9,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -57,8 +58,7 @@ class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://api.mistral.ai/v1" if api_base is None else api_base.rstrip("/") - return f"{api_base}/audio/transcriptions" + return join_api_base("https://api.mistral.ai/v1" if api_base is None else api_base, "audio/transcriptions") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return MistralAudioTranscriptionException( diff --git a/litellm/llms/mistral/batches/transformation.py b/litellm/llms/mistral/batches/transformation.py index 86de2836048..18f5df9b6e9 100644 --- a/litellm/llms/mistral/batches/transformation.py +++ b/litellm/llms/mistral/batches/transformation.py @@ -18,7 +18,7 @@ from openai.types.batch_error import BatchError from pydantic import ConfigDict from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.base import LiteLLMBaseModel @@ -156,7 +156,7 @@ class MistralBatchesConfig(BaseBatchesConfig): litellm_params: Mapping[str, object], data: CreateBatchRequest, ) -> str: - return f"{get_mistral_api_base(api_base)}/v1/batch/jobs" + return join_api_base(get_mistral_api_base(api_base), "v1/batch/jobs") def transform_create_batch_request( self, @@ -199,7 +199,10 @@ class MistralBatchesConfig(BaseBatchesConfig): api_key: Final = litellm_params.get("api_key") request: Final = MistralPresignedRequest( method="GET", - url=f"{get_mistral_api_base(api_base if isinstance(api_base, str) else None)}/v1/batch/jobs/{encoded_batch_id}", + url=join_api_base( + get_mistral_api_base(api_base if isinstance(api_base, str) else None), + f"v1/batch/jobs/{encoded_batch_id}", + ), headers=get_mistral_auth_headers(_NO_HEADERS, api_key if isinstance(api_key, str) else None), ) return dict(request) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index b8a6cb1beb9..2c0acacff23 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_internal_message_fields, strip_none_values_from_message, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, @@ -220,8 +221,7 @@ class MistralConfig(OpenAIGPTConfig): ) # if api_base does not end with /v1 we add it - if api_base is not None and not api_base.endswith("/v1"): # Mistral always needs a /v1 at the end - api_base = api_base + "/v1" + api_base = ensure_api_route(api_base, "v1") dynamic_api_key: Final = ( api_key or get_secret_str("MISTRAL_AZURE_API_KEY") # for Azure AI Mistral diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 86bce763917..908567e88fb 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -18,7 +18,7 @@ from pydantic import ConfigDict from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj from litellm.types.llms.base import LiteLLMBaseModel @@ -130,11 +130,11 @@ class MistralFilesConfig(BaseFilesConfig): litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: - return f"{get_mistral_api_base(api_base)}/v1/files" + return join_api_base(get_mistral_api_base(api_base), "v1/files") def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" + return join_api_base(_api_base_from(litellm_params), f"v1/files/{encoded_file_id}{suffix}") def get_error_class( self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers @@ -233,7 +233,7 @@ class MistralFilesConfig(BaseFilesConfig): optional_params: Mapping[str, object], litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature - url: Final = f"{_api_base_from(litellm_params)}/v1/files" + url: Final = join_api_base(_api_base_from(litellm_params), "v1/files") if not purpose: return url, _NO_QUERY_PARAMS return url, {"purpose": _to_mistral_purpose(purpose)} diff --git a/litellm/llms/modelscope/chat/transformation.py b/litellm/llms/modelscope/chat/transformation.py index c54437b9a21..ec48b344bd9 100644 --- a/litellm/llms/modelscope/chat/transformation.py +++ b/litellm/llms/modelscope/chat/transformation.py @@ -7,6 +7,7 @@ from typing import Final, Literal, cast, overload from typing_extensions import override +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -86,9 +87,4 @@ class ModelScopeChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default ModelScope /chat/completions endpoint. """ - clean_api_base: Final = (api_base or self.DEFAULT_BASE_URL).rstrip("/") - - if not clean_api_base.endswith("/chat/completions"): - return f"{clean_api_base}/chat/completions" - - return clean_api_base + return ensure_api_route(api_base or self.DEFAULT_BASE_URL, "chat/completions") diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index 08cdbacdad0..c308de92d1d 100644 --- a/litellm/llms/modelscope/image_generation/transformation.py +++ b/litellm/llms/modelscope/image_generation/transformation.py @@ -13,6 +13,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from typing_extensions import override +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -32,6 +33,7 @@ else: LiteLLMLoggingObj = object _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) @@ -90,11 +92,10 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the ModelScope image generation API request. """ - base_url: str = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL # Return the images endpoint - return f"{base_url}/images/generations" + return join_api_base(base_url, "images/generations") @override def validate_environment( @@ -173,7 +174,7 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): {"created": timestamp, "data": [{"url": "..."}]} """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing ModelScope response: {e}", diff --git a/litellm/llms/mongodb/vector_stores/transformation.py b/litellm/llms/mongodb/vector_stores/transformation.py index 0d931ee78cb..7bf9d8b8441 100644 --- a/litellm/llms/mongodb/vector_stores/transformation.py +++ b/litellm/llms/mongodb/vector_stores/transformation.py @@ -9,6 +9,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError from litellm.exceptions import AuthenticationError, BadRequestError, ServiceUnavailableError, Timeout +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, @@ -283,7 +284,7 @@ class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): raise config_error("The embedding model must return a non-empty, finite query vector.") limit: Final = cls._limit(optional_params) return ( - f"{api_base}/v1/vector_stores/{quote(vector_store_id, safe='')}/search", + join_api_base(api_base, f"v1/vector_stores/{quote(vector_store_id, safe='')}/search"), { "query": query_text, "query_vector": tuple(vector), diff --git a/litellm/llms/moonshot/chat/transformation.py b/litellm/llms/moonshot/chat/transformation.py index 93496e775db..698db8b7c25 100644 --- a/litellm/llms/moonshot/chat/transformation.py +++ b/litellm/llms/moonshot/chat/transformation.py @@ -9,6 +9,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_reasoning @@ -93,12 +94,7 @@ class MoonshotChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default Moonshot AI /chat/completions endpoint. """ - clean_api_base: Final = (api_base or "https://api.moonshot.ai/v1").rstrip("/") - - if not clean_api_base.endswith("/chat/completions"): - return f"{clean_api_base}/chat/completions" - - return clean_api_base + return ensure_api_route(api_base or "https://api.moonshot.ai/v1", "chat/completions") def get_supported_openai_params(self, model: str) -> list: """ diff --git a/litellm/llms/nadir/chat/transformation.py b/litellm/llms/nadir/chat/transformation.py index d22b1a51bab..c8bde7d9fad 100644 --- a/litellm/llms/nadir/chat/transformation.py +++ b/litellm/llms/nadir/chat/transformation.py @@ -2,8 +2,9 @@ import math from typing import Final import httpx +from pydantic import TypeAdapter -from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_response_cost_in_hidden_params from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -19,13 +20,17 @@ _SUPPORTED_OPENAI_PARAMS: Final = ( "temperature", "top_p", ) +_DECODED_JSON: Final = TypeAdapter(object) def _reported_cost_usd(raw_response: httpx.Response) -> float | None: try: - cost: Final = raw_response.json()["nadir_metadata"]["cost"]["total_cost_usd"] + body: Final = as_str_mapping(_DECODED_JSON.validate_python(raw_response.json())) except (ValueError, KeyError, TypeError): return None + metadata: Final = as_str_mapping(body.get("nadir_metadata")) if body is not None else None + cost_block: Final = as_str_mapping(metadata.get("cost")) if metadata is not None else None + cost: Final = cost_block.get("total_cost_usd") if cost_block is not None else None if isinstance(cost, bool) or not isinstance(cost, (int, float)): return None if not math.isfinite(cost) or cost < 0: diff --git a/litellm/llms/nimble/search/transformation.py b/litellm/llms/nimble/search/transformation.py index 30cfc38eb87..b4d9be909c3 100644 --- a/litellm/llms/nimble/search/transformation.py +++ b/litellm/llms/nimble/search/transformation.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Final import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, @@ -124,10 +125,7 @@ class NimbleSearchConfig(BaseSearchConfig): data: dict[str, object] | list[dict[str, object]] | None = None, # mutable-ok: base signature **kwargs: object, # kwargs-ok: BaseSearchConfig.get_complete_url signature ) -> str: - resolved_base: Final = (api_base or get_secret_str("NIMBLE_API_BASE") or self.NIMBLE_API_BASE).rstrip("/") - if resolved_base.endswith("/search"): - return resolved_base - return f"{resolved_base}/search" + return ensure_api_route(api_base or get_secret_str("NIMBLE_API_BASE") or self.NIMBLE_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/nvidia_nim/embed.py b/litellm/llms/nvidia_nim/embed.py index 5a0105a3776..d2ef446299e 100644 --- a/litellm/llms/nvidia_nim/embed.py +++ b/litellm/llms/nvidia_nim/embed.py @@ -9,6 +9,7 @@ API calling is done using the OpenAI SDK with an api_base """ import types +from collections.abc import Mapping from typing import Final @@ -32,16 +33,17 @@ class NvidiaNimEmbeddingConfig: input_type: str | None = None, truncate: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 177d7883feb..39391be0fdb 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -11,6 +11,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig from litellm.types.rerank import RerankResponse @@ -76,7 +77,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): api_base = api_base.removesuffix("/v1") - return f"{api_base}/v1/ranking" + return join_api_base(api_base, "v1/ranking") def map_cohere_rerank_params( self, diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index c8e1a331c8c..12377d8f854 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -8,6 +8,7 @@ from typing_extensions import Required, TypedDict import litellm from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -106,7 +107,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): # Strip nvidia_nim/ prefix from model name if present clean_model: Final = self._get_clean_model_name(model) - return f"{api_base}/v1/retrieval/{clean_model}/reranking" + return join_api_base(api_base, f"v1/retrieval/{clean_model}/reranking") def get_supported_cohere_rerank_params(self, model: str) -> list: """ diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index 748203b22d9..cafeece0a4e 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -28,7 +28,7 @@ without the optional STT extras installed. import asyncio import inspect -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol @@ -48,6 +48,8 @@ from litellm.llms.nvidia_riva.audio_transcription.transformation import ( RIVA_TARGET_NUM_CHANNELS, RIVA_TARGET_SAMPLE_RATE_HZ, NvidiaRivaAudioTranscriptionConfig, + RivaFinalResult, + RivaWordTiming, ) from litellm.llms.nvidia_riva.common_utils import ( NvidiaRivaException, @@ -535,14 +537,14 @@ class NvidiaRivaAudioTranscription: yield chunk @staticmethod - def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> list[dict[str, object]]: + def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> Sequence[RivaFinalResult]: """ Walk the gRPC responses, ignore empty / non-final chunks, and return a list of normalized final-result dicts. Matching the user's note: the ``id`` blocks with no ``results`` are streaming heartbeats and must be skipped. """ - final_results: Final[list[dict[str, object]]] = [] + final_results: Final[list[RivaFinalResult]] = [] for response in stream: results = getattr(response, "results", None) or [] for result in results: @@ -554,16 +556,19 @@ class NvidiaRivaAudioTranscription: top = alternatives[0] transcript = getattr(top, "transcript", "") or "" words_proto = getattr(top, "words", None) or [] - words = [] - for word in words_proto: - words.append( - { - "word": getattr(word, "word", ""), - "start_time_ms": int(getattr(word, "start_time", 0) or 0), - "end_time_ms": int(getattr(word, "end_time", 0) or 0), - } - ) - final_results.append({"transcript": transcript, "words": words}) + final_results.append( + { + "transcript": transcript, + "words": [ + RivaWordTiming( + word=getattr(word, "word", ""), + start_time_ms=int(getattr(word, "start_time", 0) or 0), + end_time_ms=int(getattr(word, "end_time", 0) or 0), + ) + for word in words_proto + ], + } + ) return final_results diff --git a/litellm/llms/nvidia_riva/audio_transcription/transformation.py b/litellm/llms/nvidia_riva/audio_transcription/transformation.py index 262dfb7bd78..5862a5d10b1 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/transformation.py +++ b/litellm/llms/nvidia_riva/audio_transcription/transformation.py @@ -11,9 +11,11 @@ dict at call time. Reference: https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-overview.html """ -from typing import Any, Final +from collections.abc import Sequence +from typing import Final, TypedDict from httpx import Headers, Response +from typing_extensions import ReadOnly from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import ( @@ -34,6 +36,17 @@ RIVA_TARGET_NUM_CHANNELS: Final = 1 RIVA_TARGET_ENCODING: Final = "LINEAR_PCM" +class RivaWordTiming(TypedDict): + word: ReadOnly[str] + start_time_ms: ReadOnly[int] + end_time_ms: ReadOnly[int] + + +class RivaFinalResult(TypedDict): + transcript: ReadOnly[str] + words: ReadOnly[Sequence[RivaWordTiming]] + + class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): """ Config for NVIDIA Riva ASR (gRPC). @@ -219,7 +232,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @staticmethod def build_transcription_response( - final_results: list[dict[str, Any]], + final_results: Sequence[RivaFinalResult], response_format: str, duration_seconds: float | None, timestamp_granularities: list[str] | None, diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index e62e5e63642..e8d37f857e6 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -21,6 +21,7 @@ from litellm._logging import verbose_logger from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -188,6 +189,7 @@ def _normalize_tool_choice(selected_params: dict) -> None: _TOOL_CHOICE_MODES: Final = frozenset({"auto", "none", "required", "any"}) _TOOLS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) _JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) def _as_json_object(value: object) -> Mapping[str, object] | None: @@ -538,7 +540,7 @@ class OCIChatConfig(BaseConfig): stream: bool | None = None, ) -> str: base: Final = get_oci_base_url(optional_params, api_base or litellm.api_base) - return f"{base}/{OCI_API_VERSION}/actions/chat" + return join_api_base(base, f"{OCI_API_VERSION}/actions/chat") def _get_optional_params(self, vendor: OCIVendors, optional_params: dict, model: str = "") -> dict: param_map: Final = ( @@ -700,7 +702,7 @@ class OCIChatConfig(BaseConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) if not isinstance(response_json, dict): raise OCIError( diff --git a/litellm/llms/oci/embed/transformation.py b/litellm/llms/oci/embed/transformation.py index edadc293771..6dee64bf19a 100644 --- a/litellm/llms/oci/embed/transformation.py +++ b/litellm/llms/oci/embed/transformation.py @@ -29,6 +29,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.oci.common_utils import ( @@ -55,6 +56,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) # OCI sends up to 96 texts per embedText request (Cohere limit). OCI_EMBED_BATCH_LIMIT: Final = 96 @@ -156,7 +158,7 @@ class OCIEmbedConfig(BaseEmbeddingConfig): stream: bool | None = None, ) -> str: base: Final = get_oci_base_url(optional_params, api_base or litellm.api_base) - return f"{base}/{OCI_API_VERSION}/actions/embedText" + return join_api_base(base, f"{OCI_API_VERSION}/actions/embedText") def sign_request( self, @@ -271,7 +273,7 @@ class OCIEmbedConfig(BaseEmbeddingConfig): ) try: - json_response: Final = raw_response.json() + json_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise OCIError( status_code=raw_response.status_code, diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 644ab187f15..f9e7071db41 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_images_from_message, extract_reasoning_content, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.ollama import ( @@ -218,10 +219,7 @@ class OllamaChatConfig(BaseConfig): Some providers need `model` in `api_base` """ - clean_api_base: Final = (api_base or "http://localhost:11434").rstrip("/") - if clean_api_base.endswith("/api/chat"): - return clean_api_base - return f"{clean_api_base}/api/chat" + return ensure_api_route(api_base or "http://localhost:11434", "api/chat") def transform_request( self, diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 284a6ffea62..f313c62f026 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -8,10 +8,12 @@ from pydantic import BaseModel, TypeAdapter from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ModelResponse, ModelResponseStream _PROVIDER_FIELDS: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) class OllamaResponseIdNormalizer: @@ -149,9 +151,9 @@ class OllamaModelInfo(BaseLLMModelInfo): names: Final[set[str]] = set() try: - resp: Final = httpx.get(f"{base}/api/tags", headers=headers) + resp: Final = httpx.get(join_api_base(base, "api/tags"), headers=headers) resp.raise_for_status() - data: Final = resp.json() + data: Final = _DECODED_JSON.validate_python(resp.json()) # Expecting a dict with a 'models' list models_list = [] if isinstance(data, dict) and "models" in data and isinstance(data["models"], list): @@ -230,7 +232,7 @@ class OllamaModelInfo(BaseLLMModelInfo): try: response: Final = module_level_client.post( - url=f"{api_base}/api/show", + url=join_api_base(api_base, "api/show"), json={"name": model}, headers=headers, ) diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 816028499e8..0c13b1cd1d3 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -10,6 +10,7 @@ from typing import Any, Final, Protocol, TypedDict from typing_extensions import NotRequired, ReadOnly import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.types.utils import EmbeddingResponse @@ -26,13 +27,6 @@ class OllamaEmbeddingResponse(TypedDict): prompt_eval_count: ReadOnly[NotRequired[int]] -def _ollama_embed_url(api_base: str) -> str: - clean_api_base: Final = api_base.rstrip("/") - if clean_api_base.endswith("/api/embed"): - return clean_api_base - return f"{clean_api_base}/api/embed" - - def _prepare_ollama_embedding_payload( model: str, prompts: list[str], optional_params: Mapping[str, object] ) -> dict[str, object]: @@ -93,13 +87,13 @@ async def ollama_aembeddings( model: str, prompts: list[str], model_response: EmbeddingResponse, - optional_params: dict, + optional_params: Mapping[str, object], logging_obj: object, encoding: TokenEncoder | None, ): data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = await litellm.module_level_aclient.post(url=_ollama_embed_url(api_base), json=data) + response: Final = await litellm.module_level_aclient.post(url=ensure_api_route(api_base, "api/embed"), json=data) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( @@ -116,14 +110,14 @@ def ollama_embeddings( api_base: str, model: str, prompts: list[str], - optional_params: dict, + optional_params: Mapping[str, object], model_response: EmbeddingResponse, logging_obj: object, encoding: TokenEncoder | None = None, ): data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = litellm.module_level_client.post(url=_ollama_embed_url(api_base), json=data) + response: Final = litellm.module_level_client.post(url=ensure_api_route(api_base, "api/embed"), json=data) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index ac41310fb61..239b4e32645 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, inline_remote_image_urls, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.base import LiteLLMBaseModel @@ -500,10 +501,7 @@ class OllamaConfig(BaseConfig): Some providers need `model` in `api_base` """ - clean_api_base: Final = (api_base or "http://localhost:11434").rstrip("/") - if clean_api_base.endswith("/api/generate"): - return clean_api_base - return f"{clean_api_base}/api/generate" + return ensure_api_route(api_base or "http://localhost:11434", "api/generate") def get_model_response_iterator( self, @@ -549,7 +547,7 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): held: Final = self.held_content self.held_content = "" try: - parsed: Final[object] = json.loads(_strip_json_code_fence(held)) + parsed: Final = _JSON_VALUE.validate_python(json.loads(_strip_json_code_fence(held))) except json.JSONDecodeError: return GenericStreamingChunk(text=held, is_finished=True, finish_reason="stop", usage=usage) if isinstance(parsed, dict) and "name" in parsed and "arguments" in parsed: diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index e21a165d757..8879c20c5eb 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -3,6 +3,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import get_httpx_client from litellm.utils import EmbeddingResponse, ModelResponse, Usage @@ -49,7 +50,7 @@ def completion( ) model = model - completion_url = completion_url + "/v1/chat/completions" + completion_url = join_api_base(completion_url, "v1/chat/completions") data: Final = oobabooga_config.transform_request( model=model, messages=messages, @@ -103,7 +104,7 @@ def embedding( if model.startswith(("http://", "https://")): embeddings_url = model elif api_base: - embeddings_url = f"{api_base}/v1/embeddings" + embeddings_url = join_api_base(api_base, "v1/embeddings") else: raise OobaboogaError( status_code=404, diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 4d7ef833181..b4c92848752 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, convert_url_to_base64, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -807,7 +808,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): if endpoint in api_base: return api_base - return f"{api_base}/{endpoint}" + return join_api_base(api_base, endpoint) def validate_environment( self, diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index e472661efb9..12ead582975 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -23,13 +23,14 @@ from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta from openai.types.completion_usage import CompletionUsage -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter if TYPE_CHECKING: from aiohttp import ClientSession import litellm from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.litellm_core_utils.token_counter import token_counter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -134,6 +135,9 @@ class OpenAIError(BaseLLMException): ) +_DECODED_JSON: Final = TypeAdapter(object) + + ####### Error Handling Utils for OpenAI API ####################### ################################################################### def drop_params_from_unprocessable_entity_error( @@ -152,19 +156,20 @@ def drop_params_from_unprocessable_entity_error( """ invalid_params: Final[list[str]] = [] if isinstance(e, httpx.HTTPStatusError): - error_json: Final = e.response.json() + error_json: Final = decoded_json_object(_DECODED_JSON.validate_python(e.response.json())) error_message: Final = error_json.get("error", {}) error_body = error_message else: error_body = e.body - if error_body is not None and isinstance(error_body, dict) and error_body.get("message"): - message = error_body.get("message", {}) + error_mapping: Final = as_str_mapping(error_body) + if isinstance(error_mapping, dict) and error_mapping.get("message"): + message = error_mapping.get("message", {}) if isinstance(message, str): try: - message = json.loads(message) + message = _DECODED_JSON.validate_python(json.loads(message)) except json.JSONDecodeError: message = {"detail": message} - detail: Final = message.get("detail") + detail: Final = decoded_json_object(message).get("detail") if isinstance(detail, list) and len(detail) > 0 and isinstance(detail[0], dict): for error_dict in detail: diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 6cb2396b85d..5d03585189d 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -9,7 +9,7 @@ from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.containers.main import ( ContainerCreateOptionalRequestParams, @@ -130,7 +130,7 @@ class OpenAIContainerConfig(BaseContainerConfig): or "https://api.openai.com/v1" ) - return f"{api_base.rstrip('/')}/containers" + return join_api_base(api_base, "containers") def transform_container_create_request( self, diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 56f348881de..25d403a0dff 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -311,6 +311,7 @@ def _openai_answer(question: DecisionsIRQuestion, answer: DecisionsIRAnswer) -> def ir_to_openai_response( response: DecisionsIRResponse, request: DecisionsIRRequest, requested_model: str ) -> OpenAIDecisionResponse: + usage: Final = response.usage or DecisionsIRUsage() return OpenAIDecisionResponse( model=response.model or requested_model, answers=tuple( @@ -318,14 +319,14 @@ def ir_to_openai_response( for question, answer in zip(request.questions, response.answers, strict=True) ), usage=OpenAIDecisionUsage( - input_tokens=response.usage.input_tokens, + input_tokens=usage.input_tokens, input_tokens_details=OpenAIDecisionInputTokensDetails( - cached_tokens=response.usage.cached_tokens, - cache_write_tokens=response.usage.cache_write_tokens, + cached_tokens=usage.cached_tokens, + cache_write_tokens=usage.cache_write_tokens, ), - output_tokens=response.usage.output_tokens, - output_tokens_details=OpenAIDecisionOutputTokensDetails(reasoning_tokens=response.usage.reasoning_tokens), - total_tokens=response.usage.input_tokens + response.usage.output_tokens, + output_tokens=usage.output_tokens, + output_tokens_details=OpenAIDecisionOutputTokensDetails(reasoning_tokens=usage.reasoning_tokens), + total_tokens=usage.input_tokens + usage.output_tokens, ), ) diff --git a/litellm/llms/openai/evals/transformation.py b/litellm/llms/openai/evals/transformation.py index 24de7f6d479..73d211e1613 100644 --- a/litellm/llms/openai/evals/transformation.py +++ b/litellm/llms/openai/evals/transformation.py @@ -2,13 +2,13 @@ OpenAI Evals API configuration and transformations """ -from collections.abc import Mapping from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.evals.transformation import ( BaseEvalsAPIConfig, LiteLLMLoggingObj, @@ -32,9 +32,11 @@ from litellm.types.llms.openai_evals import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +_DECODED_JSON: Final = TypeAdapter(object) -def _parsed_response_json(raw_response: httpx.Response) -> Mapping[str, object]: - return raw_response.json() + +def _parsed_response_json(raw_response: httpx.Response) -> object: + return _DECODED_JSON.validate_python(raw_response.json()) class OpenAIEvalsConfig(BaseEvalsAPIConfig): @@ -67,8 +69,8 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): if eval_id: encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") - return f"{api_base}/v1/evals/{encoded_eval_id}" - return f"{api_base}/v1/{endpoint}" + return join_api_base(api_base, f"v1/evals/{encoded_eval_id}") + return join_api_base(api_base, f"v1/{endpoint}") def transform_create_eval_request( self, @@ -226,7 +228,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> tuple[str, dict, dict]: """Transform cancel eval request for OpenAI""" - url: Final = f"{self.get_complete_url(api_base=api_base, endpoint='evals', eval_id=eval_id)}/cancel" + url: Final = join_api_base( + self.get_complete_url(api_base=api_base, endpoint="evals", eval_id=eval_id), "cancel" + ) # Empty body for cancel request request_body: Final[dict[str, object]] = {} @@ -260,7 +264,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base = litellm_params.api_base encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs") # Build request body request_body: Final = {k: v for k, v in create_request.items() if v is not None} @@ -293,7 +297,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base = litellm_params.api_base encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs") # Build query parameters query_params: Final[dict[str, object]] = {} @@ -335,7 +339,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): """Transform get run request for OpenAI""" encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs/{encoded_run_id}") verbose_logger.debug("Get run request - URL: %s", url) @@ -363,7 +367,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): """Transform cancel run request for OpenAI""" encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}/cancel" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs/{encoded_run_id}/cancel") # Empty body for cancel request request_body: Final[dict[str, object]] = {} @@ -394,7 +398,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): """Transform delete run request for OpenAI""" encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs/{encoded_run_id}") # Empty body for delete request request_body: Final[dict[str, object]] = {} diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index a9e5b1b9f34..9def833052b 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -6,6 +6,7 @@ from httpx._types import RequestFiles import litellm from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ( @@ -184,7 +185,4 @@ class OpenAIImageEditConfig(BaseImageEditConfig): or "https://api.openai.com/v1" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/images/edits" + return join_api_base(api_base, "images/edits") diff --git a/litellm/llms/openai/realtime/http_transformation.py b/litellm/llms/openai/realtime/http_transformation.py index 61dbf20397f..49d33582c6f 100644 --- a/litellm/llms/openai/realtime/http_transformation.py +++ b/litellm/llms/openai/realtime/http_transformation.py @@ -1,6 +1,7 @@ """OpenAI realtime HTTP transformation config (client_secrets + realtime_calls).""" import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig from litellm.secret_managers.main import get_secret_str @@ -15,17 +16,17 @@ class OpenAIRealtimeHTTPConfig(BaseRealtimeHTTPConfig): def get_complete_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: base = self.get_api_base(api_base).rstrip("/") base = base.removesuffix("/v1") - return f"{base}/v1/realtime/client_secrets" + return join_api_base(base, "v1/realtime/client_secrets") def get_realtime_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: base = self.get_api_base(api_base).rstrip("/") base = base.removesuffix("/v1") - return f"{base}/v1/realtime/calls" + return join_api_base(base, "v1/realtime/calls") def get_transcription_session_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: base = self.get_api_base(api_base).rstrip("/") base = base.removesuffix("/v1") - return f"{base}/v1/realtime/transcription_sessions" + return join_api_base(base, "v1/realtime/transcription_sessions") def validate_environment( self, diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 437f19a4dc6..deee502c9f2 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -20,7 +20,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.transformation impor LiteLLMAnthropicToResponsesAPIAdapter, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials +from litellm.llms.openai.common_utils import OpenAIError, decoded_json_object, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, ) @@ -244,7 +244,7 @@ class OpenAITokenCounter(BaseTokenCounter): if result is not None: return TokenCountResponse.model_validate( { - "total_tokens": result.get("input_tokens", 0), + "total_tokens": decoded_json_object(result).get("input_tokens", 0), "request_model": request_model, "model_used": model_to_use, "tokenizer_type": "openai_api", diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 6d882e0eeb1..0f885146e36 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -6,13 +6,15 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp import copy from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal, TypeAlias +from typing import Final, Literal, Protocol, TypeAlias from openai.types.responses.input_token_count_params import InputTokenCountParams from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes from litellm.litellm_core_utils.tool_fields import DropPolicy +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.pass_through.convertible_blocks import malformed_tool_result_text_error from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( @@ -60,6 +62,26 @@ class _AnthropicCountMessage(TypedDict): _ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...]) + + +class _FieldReader(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +class _FunctionFields(_FieldReader, Protocol): + def __contains__(self, key: object, /) -> bool: ... + + def __getitem__(self, key: str, /) -> object: ... + + +def _field_reader(value: object) -> _FieldReader: + return require_attributes(value, _FieldReader, "get") + + +def _function_fields(value: object) -> _FunctionFields: + return require_attributes(value, _FunctionFields, "get", "__contains__", "__getitem__") + + _ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema) @@ -148,9 +170,8 @@ class OpenAICountTokensConfig: """ def get_openai_count_tokens_endpoint(self, api_base: str | None = None) -> str: - base = api_base or "https://api.openai.com/v1" - base = base.rstrip("/") - return f"{base}/responses/input_tokens" + base: Final = api_base or "https://api.openai.com/v1" + return join_api_base(base, "responses/input_tokens") def transform_request_to_count_tokens( self, @@ -196,7 +217,7 @@ class OpenAICountTokensConfig: @staticmethod def _transform_tools_for_responses_api( - tools: list[dict[str, Any]], + tools: list[dict[str, object]], ) -> list[dict[str, object]]: """ Transform OpenAI chat and Anthropic tools to Responses API tools format. @@ -208,7 +229,7 @@ class OpenAICountTokensConfig: transformed: Final[list[dict[str, object]]] = [] for tool in tools: if tool.get("type") == "function" and "function" in tool: - func = tool["function"] + func = _function_fields(tool["function"]) item: dict[str, object] = { "type": "function", "name": func.get("name", ""), @@ -235,7 +256,7 @@ class OpenAICountTokensConfig: @staticmethod def messages_to_responses_input( - messages: list[dict[str, Any]], + messages: Sequence[Mapping[str, object]], ) -> tuple[list[dict[str, object]], str | None]: """ Convert standard chat messages format to OpenAI Responses API input format. @@ -276,14 +297,17 @@ class OpenAICountTokensConfig: if content: input_items.append({"role": "assistant", "content": content}) if tool_calls: - for tc in tool_calls: - func = tc.get("function", {}) + for tc in iter_items(tool_calls): + tool_call = _field_reader(tc) + function = tool_call.get("function", {}) + call_id = tool_call.get("id", "") + function_fields = _field_reader(function) input_items.append( { "type": "function_call", - "call_id": tc.get("id", ""), - "name": func.get("name", ""), - "arguments": func.get("arguments", ""), + "call_id": call_id, + "name": function_fields.get("name", ""), + "arguments": function_fields.get("arguments", ""), } ) elif not content: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 84ff790ae48..d37a12a91d6 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -36,7 +36,7 @@ from collections.abc import Mapping, MutableSequence, Sequence from dataclasses import dataclass from itertools import accumulate, chain, repeat from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, Union, cast from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -47,6 +47,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i OpenAiResponsesToChatCompletionStreamIterator, tool_call_dict_from_output_item, ) +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -65,6 +66,11 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( unappliable_request_rewrite, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + decode_compaction_item, + with_compaction_summary, +) from litellm.responses.litellm_completion_transformation.custom_tools import without_malformed_tool_flags from litellm.responses.litellm_completion_transformation.tool_search import is_server_tool_search_tool from litellm.responses.litellm_completion_transformation.transformation import ( @@ -110,6 +116,16 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +_COMPACTION_SUMMARY_POSITION: Final = -1 + + +def _output_compaction_summary(item: object) -> str | None: + if stream_item_field(item, "type") != "compaction": + return None + block: Final = decode_compaction_item(stream_item_field(item, "encrypted_content")) + return None if block is None else block.summary + + class _ToolCallShape(NamedTuple): name: str | None arguments: str @@ -222,6 +238,10 @@ class ResponsesStreamChunk(TypedDict, total=False): content_index: ReadOnly[int] +class _StreamEventFields(Protocol): + def get(self, key: str, /) -> object: ... + + _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( { ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, @@ -266,7 +286,17 @@ def _item_rewrite_field(item: Mapping[str, object]) -> str | None: return _PATCHABLE_ITEM_FIELDS.get(item_type) +def _rewritten_compaction_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + encrypted_content: Final = item.get("encrypted_content") + summary: Final = rewritten.get("content") if isinstance(rewritten, Mapping) else None + if not isinstance(encrypted_content, str) or not isinstance(summary, str) or not summary: + return None + return {**item, "encrypted_content": with_compaction_summary(encrypted_content, summary)} + + def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + if compaction_summary(item) is not None: + return _rewritten_compaction_item(item, rewritten) field: Final = _item_rewrite_field(item) if field is None or not isinstance(rewritten, Mapping): return None @@ -804,6 +834,11 @@ class OpenAIResponsesHandler(BaseTranslation): Override this method to customize text/image extraction logic. """ + summary: Final = compaction_summary(message) + if summary is not None: + texts_to_check.append(summary) + task_mappings.append((msg_idx, None)) + return content: Final = message.get("content", None) if content is None: return @@ -849,6 +884,12 @@ class OpenAIResponsesHandler(BaseTranslation): msg_idx = cast(int, mapping[0]) content_idx_optional = cast(int | None, mapping[1]) + if compaction_summary(messages[msg_idx]) is not None: + messages[msg_idx]["encrypted_content"] = with_compaction_summary( # rebind-ok: in-place input rewrite + messages[msg_idx]["encrypted_content"], guardrail_response + ) + continue + content = messages[msg_idx].get("content", None) if content is None: continue @@ -1116,13 +1157,17 @@ class OpenAIResponsesHandler(BaseTranslation): # delivered, so a truncated buffer still reports text undeliverable. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": + finished_summary: Final = _output_compaction_summary(stream_item_field(final_chunk, "item")) model_response_stream: Final = ( OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) ) tool_calls: Final = model_response_stream.choices[0].delta.tool_calls - if tool_calls: + if tool_calls or finished_summary is not None: inputs = GenericGuardrailAPIInputs() - inputs["tool_calls"] = cast(list[ChatCompletionToolCallChunk], tool_calls) + if finished_summary is not None: + inputs["texts"] = [finished_summary] + if tool_calls: + inputs["tool_calls"] = cast(list[ChatCompletionToolCallChunk], tool_calls) if hasattr(model_response_stream, "model") and model_response_stream.model: inputs["model"] = model_response_stream.model await guardrail_to_apply.apply_guardrail( @@ -1214,9 +1259,17 @@ class OpenAIResponsesHandler(BaseTranslation): else: setattr(event, field, value) + @staticmethod + def _write_compaction_summary(item: object, summary: str) -> None: + encrypted_content: Final = stream_item_field(item, "encrypted_content") + if isinstance(encrypted_content, str): + OpenAIResponsesHandler._write_event_field( + item, "encrypted_content", with_compaction_summary(encrypted_content, summary) + ) + def _sync_stream_events_with_rewrites( self, - stream_events: Sequence[Any], + stream_events: Sequence[object], rewrites_by_position: Mapping[tuple[int, int], str], ) -> None: """Sync pre-completion stream events with the rewritten completed @@ -1230,7 +1283,7 @@ class OpenAIResponsesHandler(BaseTranslation): {position: chain((rewritten,), repeat("")) for position, rewritten in rewrites_by_position.items()} ) for event in stream_events: - if not (isinstance(event, dict) or hasattr(event, "get")): + if not has_attributes(event, _StreamEventFields, "get"): continue event_type = event.get("type") output_index = event.get("output_index") @@ -1256,11 +1309,13 @@ class OpenAIResponsesHandler(BaseTranslation): output_index: int, rewrites_by_position: Mapping[tuple[int, int], str], ) -> None: - content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None) - if not isinstance(content, list): + summary_rewrite: Final = rewrites_by_position.get((output_index, _COMPACTION_SUMMARY_POSITION)) + if summary_rewrite is not None: + OpenAIResponsesHandler._write_compaction_summary(item, summary_rewrite) return + content: Final = stream_item_items(item, "content") for (item_idx, content_idx), rewritten in rewrites_by_position.items(): - if item_idx != output_index or content_idx >= len(content): + if item_idx != output_index or not 0 <= content_idx < len(content): continue OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) @@ -1481,13 +1536,13 @@ class OpenAIResponsesHandler(BaseTranslation): message_items: Final = tuple( item for item in output_items if stream_item_field(item, "type") not in _OUTPUT_TOOL_CALL_ITEM_TYPES ) + message_parts: Final = chain.from_iterable(stream_item_items(item, "content") for item in message_items) + part_texts: Final = tuple( + text for part in message_parts if isinstance(text := stream_item_field(part, "text"), str) and text + ) + summaries: Final = tuple(summary for item in output_items if (summary := _output_compaction_summary(item))) return StreamingScanKey( - texts=tuple( - text - for item in message_items - for part in stream_item_items(item, "content") - if isinstance(text := stream_item_field(part, "text"), str) and text - ), + texts=(*part_texts, *summaries), tool_calls=tuple( stream_item_fingerprint(item) for item in output_items @@ -1610,6 +1665,12 @@ class OpenAIResponsesHandler(BaseTranslation): ) return + summary: Final = _output_compaction_summary(output_item) + if summary is not None: + texts_to_check.append(summary) + task_mappings.append((output_idx, _COMPACTION_SUMMARY_POSITION)) + return + # Handle both GenericResponseOutputItem and dict content: list[OutputText] | list[dict] | None = None if isinstance(output_item, BaseModel): @@ -1677,6 +1738,10 @@ class OpenAIResponsesHandler(BaseTranslation): output_item = response_output[output_idx] + if content_idx == _COMPACTION_SUMMARY_POSITION: + self._write_compaction_summary(output_item, guardrail_response) + continue + # Handle both GenericResponseOutputItem, BaseModel, and dict if isinstance(output_item, GenericResponseOutputItem): if output_item.content and content_idx < len(output_item.content): diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 8f29a71d63b..b90f6aa01da 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -23,7 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.openai.chat.gpt_5_transformation import ( is_gpt_reasoning_series_name, @@ -881,10 +881,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): or "https://api.openai.com/v1" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def transform_streaming_response( self, @@ -1111,7 +1108,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): order: Literal["asc", "desc"] = "desc", ) -> tuple[str, dict]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") - url: Final = f"{api_base}/{encoded_response_id}/input_items" + url: Final = join_api_base(api_base, f"{encoded_response_id}/input_items") params: Final[dict[str, object]] = {} if after is not None: params["after"] = after @@ -1152,7 +1149,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): - POST /v1/responses/{response_id}/cancel """ encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") - url: Final = f"{api_base}/{encoded_response_id}/cancel" + url: Final = join_api_base(api_base, f"{encoded_response_id}/cancel") data: Final[dict] = {} return url, data diff --git a/litellm/llms/openai/transcriptions/whisper_transformation.py b/litellm/llms/openai/transcriptions/whisper_transformation.py index 8f96f688189..eb4f1b6c9ce 100644 --- a/litellm/llms/openai/transcriptions/whisper_transformation.py +++ b/litellm/llms/openai/transcriptions/whisper_transformation.py @@ -4,6 +4,7 @@ from typing import Final from httpx import Headers, Response +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -46,9 +47,9 @@ class OpenAIWhisperAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # if endswith "/v1" if api_base and api_base.endswith("/v1"): - api_base = f"{api_base}/audio/transcriptions" + api_base = join_api_base(api_base, "audio/transcriptions") else: - api_base = f"{api_base}/v1/audio/transcriptions" + api_base = join_api_base(api_base, "v1/audio/transcriptions") return api_base or "" diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index 7f65ade224e..bd8145bb460 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -4,7 +4,7 @@ from typing import Final, cast import httpx import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) @@ -92,16 +92,15 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): vector_store_id: str, litellm_params: Mapping[str, object], ) -> str: - base_url = ( + base_url: Final = ( api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - base_url = base_url.rstrip("/") encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - return f"{base_url}/vector_stores/{encoded_vector_store_id}/files" + return join_api_base(base_url, f"vector_stores/{encoded_vector_store_id}/files") def transform_create_vector_store_file_request( self, @@ -187,7 +186,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): api_base: str, ) -> tuple[str, dict[str, object]]: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{api_base}/{encoded_file_id}/content", {} + return join_api_base(api_base, f"{encoded_file_id}/content"), {} def transform_retrieve_vector_store_file_content_response( self, diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 46105a66be3..d108802f7f9 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.openai.workload_identity import resolve_openai_bearer_token from litellm.secret_managers.main import get_secret_str @@ -88,10 +88,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): or "https://api.openai.com/v1" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/vector_stores" + return join_api_base(api_base, "vector_stores") def transform_search_vector_store_request( self, @@ -105,7 +102,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") base_url, query_separator, query_string = api_base.partition("?") - url: Final = f"{base_url}/{encoded_vector_store_id}/search{query_separator}{query_string}" + url: Final = f"{join_api_base(base_url, f'{encoded_vector_store_id}/search')}{query_separator}{query_string}" typed_request_body: Final = VectorStoreSearchRequest( query=query, filters=vector_store_search_optional_params.get("filters", None), diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index d422e51f649..0560e1c2e3a 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -8,7 +8,7 @@ from urllib.parse import quote import httpx from httpx._types import FileContent, FileTypes, RequestFiles -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.openai.image_edit.transformation import ImageEditRequestUtils from litellm.llms.openai.workload_identity import resolve_openai_bearer_token @@ -98,7 +98,7 @@ class OpenAIVideoConfig(BaseVideoConfig): if api_base is None: api_base = "https://api.openai.com/v1" - return f"{api_base.rstrip('/')}/videos" + return join_api_base(api_base, "videos") def use_multipart_form_data(self) -> bool: return True @@ -214,7 +214,7 @@ class OpenAIVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for video content download - url = f"{api_base.rstrip('/')}/{encoded_video_id}/content" + url = join_api_base(api_base, f"{encoded_video_id}/content") if variant is not None: # Encode the user-controlled ``variant`` so a value like # ``thumbnail&extra=1`` cannot inject additional query params @@ -246,7 +246,7 @@ class OpenAIVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for video remix - url: Final = f"{api_base.rstrip('/')}/{encoded_video_id}/remix" + url: Final = join_api_base(api_base, f"{encoded_video_id}/remix") # Prepare the request data data: Final[dict[str, object]] = {"prompt": prompt} @@ -465,7 +465,7 @@ class OpenAIVideoConfig(BaseVideoConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> tuple[str, list]: - url: Final = f"{api_base.rstrip('/')}/characters" + url: Final = join_api_base(api_base, "characters") files_list: Final[list[tuple[str, FileTypes]]] = [("name", (None, name))] self._add_video_to_files(files_list, video, "video") return url, files_list @@ -486,7 +486,7 @@ class OpenAIVideoConfig(BaseVideoConfig): ) -> tuple[str, dict]: original_character_id: Final = extract_original_character_id(character_id) encoded_character_id: Final = encode_url_path_segment(original_character_id, field_name="character_id") - url: Final = f"{api_base.rstrip('/')}/characters/{encoded_character_id}" + url: Final = join_api_base(api_base, f"characters/{encoded_character_id}") return url, {} def transform_video_get_character_response( @@ -507,7 +507,7 @@ class OpenAIVideoConfig(BaseVideoConfig): extra_body: dict[str, object] | None = None, prefetched_source_data: dict[str, object] | None = None, ) -> tuple[str, Mapping[str, object], RequestFiles | None]: - url: Final = f"{api_base.rstrip('/')}/edits" + url: Final = join_api_base(api_base, "edits") if video_file is not None: files: Final[RequestFiles] = (self._video_file_tuple(video_file, "video"),) @@ -547,7 +547,7 @@ class OpenAIVideoConfig(BaseVideoConfig): extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: original_video_id: Final = extract_original_video_id(video_id) - url: Final = f"{api_base.rstrip('/')}/extensions" + url: Final = join_api_base(api_base, "extensions") data: Final[dict[str, object]] = { "prompt": prompt, "seconds": seconds, diff --git a/litellm/llms/openai_like/common_utils.py b/litellm/llms/openai_like/common_utils.py index 11b85e52af5..5d3085a3901 100644 --- a/litellm/llms/openai_like/common_utils.py +++ b/litellm/llms/openai_like/common_utils.py @@ -2,6 +2,8 @@ from typing import Literal import httpx +from litellm.litellm_core_utils.url_utils import join_api_base + class OpenAILikeError(Exception): def __init__(self, status_code, message): @@ -48,7 +50,7 @@ class OpenAILikeBase: if not custom_endpoint: if endpoint_type == "chat_completions": - api_base = f"{api_base}/chat/completions" + api_base = join_api_base(api_base, "chat/completions") elif endpoint_type == "embeddings": - api_base = f"{api_base}/embeddings" + api_base = join_api_base(api_base, "embeddings") return api_base, headers diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index f0bf710b951..685561f55a5 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -9,6 +9,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.secret_managers.main import get_secret_str @@ -91,13 +92,9 @@ def create_config_class(provider: SimpleProviderConfig): if not api_base: api_base = provider.base_url - if api_base is None: - raise ValueError(f"api_base is required for provider {provider.slug}") - - if not api_base.endswith("/chat/completions"): - api_base = f"{api_base}/chat/completions" - - return api_base + if isinstance(api_base, str): + return ensure_api_route(api_base, "chat/completions") + raise ValueError(f"api_base is required for provider {provider.slug}") def get_supported_openai_params(self, model: str) -> list: """Get supported OpenAI params, excluding tool-related params for models @@ -226,11 +223,9 @@ def create_responses_config_class(provider: SimpleProviderConfig) -> "type[OpenA if not api_base: api_base = provider.base_url - if api_base is None: - raise ValueError(f"api_base is required for provider {provider.slug}") - - api_base = api_base.rstrip("/") - return f"{api_base}/responses" + if isinstance(api_base, str): + return join_api_base(api_base, "responses") + raise ValueError(f"api_base is required for provider {provider.slug}") def transform_responses_api_request( self, diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index ef501720e2d..0ab55d653e0 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -1,6 +1,7 @@ -from typing import Any, Final +from typing import Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, @@ -40,7 +41,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): self, headers: dict[str, str], model: str, - messages: list[Any], + messages: list[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, @@ -113,7 +114,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): if base.endswith("/v1/messages"): return base base = base.removesuffix("/v1") - return f"{base}/v1/messages" + return join_api_base(base, "v1/messages") class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): @@ -150,7 +151,7 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): self, headers: dict[str, str], model: str, - messages: list[Any], + messages: list[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/openai_like/responses/transformation.py b/litellm/llms/openai_like/responses/transformation.py index 9655772e953..68de4065b09 100644 --- a/litellm/llms/openai_like/responses/transformation.py +++ b/litellm/llms/openai_like/responses/transformation.py @@ -8,6 +8,7 @@ provider-specific resolution (slug, API key env var, base URL). from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -47,5 +48,4 @@ class OpenAILikeResponsesConfig(OpenAIResponsesAPIConfig): api_base = api_base or get_secret_str("OPENAI_LIKE_API_BASE") if not api_base: raise ValueError("api_base is required for openai_like provider") - api_base = api_base.rstrip("/") - return f"{api_base}/responses" + return join_api_base(api_base, "responses") diff --git a/litellm/llms/opencode/chat/__init__.py b/litellm/llms/opencode/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/opencode/chat/transformation.py b/litellm/llms/opencode/chat/transformation.py new file mode 100644 index 00000000000..dcb30b5835c --- /dev/null +++ b/litellm/llms/opencode/chat/transformation.py @@ -0,0 +1,248 @@ +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import TypeAdapter + +import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.base_llm.chat.transformation import BaseConfig +from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig +from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig +from litellm.llms.opencode.common_utils import ( + attach_go_reported_cost, + go_api_base, + go_api_key, + go_endpoint, + go_headers, + go_model_info, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + + +class _GoReportedCostMixin: + """Record usage.opencode.list_cost so response_cost reflects OpenCode's price, not the static rows.""" + + go_usage_key: str = "usage" + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "Logging", + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: object, + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + response: Final = super().transform_response( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # mixin precedes a BaseConfig subclass + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=dict(request_data), + messages=list(messages), + optional_params=dict(optional_params), + litellm_params=dict(litellm_params), + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + priced: Final = TypeAdapter(ModelResponse).validate_python(response) + attach_go_reported_cost( + priced, TypeAdapter(dict[str, object]).validate_json(raw_response.content).get(self.go_usage_key) + ) + return priced + + +class OpenCodeGoChatConfig(_GoReportedCostMixin, OpenAILikeChatConfig): + @property + def custom_llm_provider(self) -> str: + return "opencode_go" + + def _get_openai_compatible_provider_info( + self, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str | None]: + return go_api_base(api_base), go_api_key(api_key) + + def validate_environment( + self, + headers: Mapping[str, object], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: protocol handlers mutate the returned request and headers + return go_headers(headers, litellm_params, api_key, model) + + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: provider discovery requires a list + response: Final = litellm.module_level_client.get( # pyright: ignore[reportUnknownMemberType] # legacy HTTP wrapper has untyped query parameters + url=ensure_api_route( + (api_base or get_secret_str("OPENCODE_GO_API_BASE") or "https://opencode.ai/inference/go") + .rstrip("/") + .removesuffix("/v1"), + "v1/models", + ), + headers={"Authorization": f"Bearer {go_api_key(api_key)}"} if go_api_key(api_key) else {}, + ) + response.raise_for_status() + data: Final = TypeAdapter(dict[str, object]).validate_json(response.content) + models: Final = TypeAdapter(tuple[Mapping[str, object], ...]).validate_python(data["data"]) + return [f"opencode_go/{entry['id']}" for entry in models if isinstance(entry.get("id"), str)] + + +class OpenCodeGoMessagesChatConfig(_GoReportedCostMixin, AnthropicConfig): + @property + def custom_llm_provider(self) -> str: + return "opencode_go" + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + return ensure_api_route(go_api_base(api_base, model).rstrip("/").removesuffix("/v1"), "v1/messages") + + def validate_environment( + self, + headers: Mapping[str, object], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: protocol handlers mutate the returned request and headers + return go_headers({"anthropic-version": "2023-06-01", **headers}, litellm_params, api_key, model) + + def transform_request( + self, + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: protocol handlers mutate the returned request and headers + max_output: Final[int | None] = TypeAdapter[int | None](int | None).validate_python( + go_model_info(model).get("max_output_tokens") + ) + defaults: Final = {"max_tokens": max_output} if max_output else {} + raw: Final[object] = super().transform_request( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy Anthropic transformer accepts untyped mappings + model, + list(messages), + {**defaults, **optional_params}, + dict(litellm_params), + dict(headers), + ) + return TypeAdapter(dict[str, object]).validate_python(raw) + + +class OpenCodeGoGeminiChatConfig(_GoReportedCostMixin, GoogleAIStudioGeminiConfig): + go_usage_key: str = "usageMetadata" + + @property + def supports_stream_param_in_request_body(self) -> bool: + return False + + def get_model_response_iterator_for_request( + self, + streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, + sync_stream: bool, + json_mode: bool | None, + optional_params: Mapping[str, object], + logging_obj: "Logging | None" = None, + ) -> object: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ModelResponseIterator + + if logging_obj is None: + raise ValueError("OpenCode Go Gemini streaming requires the request logging context") + iterator: Final = ModelResponseIterator( + streaming_response=streaming_response, sync_stream=sync_stream, logging_obj=logging_obj + ) + return iter(iterator) if sync_stream else iterator + + @property + def custom_llm_provider(self) -> str: + return "opencode_go" + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + base: Final = ensure_api_route( + go_api_base(api_base, model).rstrip("/").removesuffix("/v1beta"), "v1beta/models" + ) + action: Final = "streamGenerateContent?alt=sse" if stream else "generateContent" + return f"{base}/{model}:{action}" + + def validate_environment( + self, + headers: Mapping[str, object] | None, + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | Mapping[str, object] | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: Gemini handler requires mutable string headers + return TypeAdapter(dict[str, str]).validate_python( + go_headers(headers or {}, litellm_params, api_key if isinstance(api_key, str) else None, model) + ) + + def transform_request( + self, + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: protocol handlers mutate the returned request and headers + from litellm.llms.vertex_ai.gemini.transformation import ( + _transform_request_body, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared legacy Gemini wire-format builder + ) + + cached_content: Final = optional_params.get("cached_content") + body: Final = _transform_request_body( + messages=list(messages), + model=model, + optional_params=dict(optional_params), + custom_llm_provider="gemini", + litellm_params=dict(litellm_params), + cached_content=cached_content if isinstance(cached_content, str) else None, + ) + return TypeAdapter(dict[str, object]).validate_python(body) + + +def go_chat_config(model: str) -> BaseConfig: + match go_endpoint(model): + case "/v1/messages": + return OpenCodeGoMessagesChatConfig() + case "/v1beta/models": + return OpenCodeGoGeminiChatConfig() + case _: + return OpenCodeGoChatConfig() diff --git a/litellm/llms/opencode/common_utils.py b/litellm/llms/opencode/common_utils.py new file mode 100644 index 00000000000..82cd957d23c --- /dev/null +++ b/litellm/llms/opencode/common_utils.py @@ -0,0 +1,80 @@ +import re +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import provider_reported_usage_cost, set_response_cost_in_hidden_params +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ModelResponse + + +def is_opencode_api_base(api_base: str | None) -> bool: + return bool( + api_base + and re.fullmatch( + r"https://opencode\.ai/(?:zen(?:/go)?|inference/(?:go/)?(?:openai|anthropic))(?:/v1)?/?", api_base + ) + ) + + +def go_model_info(model: str) -> Mapping[str, object]: + import litellm + + raw: Final[object] = litellm.model_cost.get( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy mutable catalog, validated below + f"opencode_go/{model.removeprefix('opencode_go/')}", {} + ) + return TypeAdapter(Mapping[str, object]).validate_python(raw) + + +def go_endpoint(model: str) -> str: + endpoints: Final = TypeAdapter(tuple[str, ...]).validate_python( + go_model_info(model).get("supported_endpoints") or () + ) + return next(iter(endpoints), "/v1/chat/completions") + + +def go_api_base(api_base: str | None, model: str = "") -> str: + endpoint: Final = go_endpoint(model) + family: Final = ( + "anthropic/v1" + if endpoint == "/v1/messages" + else ("google/v1beta" if endpoint == "/v1beta/models" else "openai/v1") + ) + return api_base or get_secret_str("OPENCODE_GO_API_BASE") or f"https://opencode.ai/inference/go/{family}" + + +def go_api_key(api_key: str | None) -> str | None: + return api_key or get_secret_str("OPENCODE_GO_API_KEY") or get_secret_str("OPENCODE_API_KEY") + + +def go_headers( + headers: Mapping[str, object], + params: object, + api_key: str | None, + model: str, +) -> dict[str, object]: # mutable-ok: protocol handlers mutate the returned request and headers + import litellm + + resolved_key: Final = go_api_key(api_key) + if not resolved_key: + raise litellm.AuthenticationError( + message="Set OPENCODE_API_KEY or pass api_key for OpenCode Go.", + llm_provider="opencode_go", + model=model, + ) + return add_provider_affinity_header( + {"Content-Type": "application/json", **headers, "Authorization": f"Bearer {resolved_key}"}, + params, + "x-opencode-session", + ) + + +def go_reported_cost(usage: object) -> float | None: + return provider_reported_usage_cost(usage, "opencode_go") + + +def attach_go_reported_cost(response: ModelResponse | ResponsesAPIResponse, usage: object) -> None: + set_response_cost_in_hidden_params(response, go_reported_cost(usage)) diff --git a/litellm/llms/opencode/harness/transformation.py b/litellm/llms/opencode/harness/transformation.py index 46e8b750ca5..e801b66c448 100644 --- a/litellm/llms/opencode/harness/transformation.py +++ b/litellm/llms/opencode/harness/transformation.py @@ -26,10 +26,10 @@ import itertools import json import secrets import string -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError @@ -45,6 +45,7 @@ from litellm.harness.types import ( ToolCall, ToolResult, ) +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.base_llm.harness.transformation import ( BaseCLIHarnessConfig, HarnessSessionCopy, @@ -269,10 +270,17 @@ def session_copy(exported: str, path: str, workdir: str) -> HarnessSessionCopy | ) -def _as_dict(value: object) -> Mapping[str, Any]: +def _as_dict(value: object) -> Mapping[str, object]: return value if isinstance(value, dict) else MappingProxyType({}) +def _spread_items(value: object) -> Iterable[object]: + try: + return iter_items(value) + except TypeError: + raise TypeError(f"Value after * must be an iterable, not {type(value).__name__}") from None + + def _tool_events(part: Mapping[str, object]) -> Sequence[Event]: native = str(part.get("tool") or "") call_id = str(part.get("callID") or part.get("id") or "") @@ -346,7 +354,7 @@ def build_opencode_config( token_path: str, permissions: PermissionMode, disable_tools: Sequence[str] = (), - user_config: Mapping[str, Any] | None = None, + user_config: Mapping[str, object] | None = None, instructions_path: str | None = None, skills_path: str | None = None, ) -> Mapping[str, object]: @@ -355,10 +363,16 @@ def build_opencode_config( validate_user_config(user) qualified = f"{OPENCODE_PROVIDER_ID}/{model}" extra_instructions: Final = (instructions_path,) if instructions_path else () - instructions: Final = [*(user.get("instructions") or ()), *extra_instructions] # mutable-ok: opencode config JSON + instructions: Final = [ + *_spread_items(user.get("instructions") or ()), + *extra_instructions, + ] # mutable-ok: opencode config JSON user_skills = _as_dict(user.get("skills")) extra_skills: Final = (skills_path,) if skills_path else () - skill_paths: Final = [*(user_skills.get("paths") or ()), *extra_skills] # mutable-ok: opencode config JSON + skill_paths: Final = [ + *_spread_items(user_skills.get("paths") or ()), + *extra_skills, + ] # mutable-ok: opencode config JSON options: Final = {"baseURL": base_url, "apiKey": "{file:" + token_path + "}"} # mutable-ok: opencode config JSON models: Final[dict[str, Mapping[str, object]]] = {model: {}} # mutable-ok: opencode config JSON provider: Final = { # mutable-ok: opencode config JSON diff --git a/litellm/llms/opencode/messages/__init__.py b/litellm/llms/opencode/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/opencode/messages/transformation.py b/litellm/llms/opencode/messages/transformation.py new file mode 100644 index 00000000000..b08a6df47aa --- /dev/null +++ b/litellm/llms/opencode/messages/transformation.py @@ -0,0 +1,47 @@ +from collections.abc import Mapping, Sequence + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.anthropic.pass_through.messages.transformation import AnthropicMessagesConfig +from litellm.llms.opencode.common_utils import go_api_base, go_headers + + +class OpenCodeGoMessagesConfig(AnthropicMessagesConfig): + @property + def custom_llm_provider(self) -> str: + return "opencode_go" + + def should_strip_billing_metadata(self) -> bool: + return True + + def validate_anthropic_messages_environment( + self, + headers: Mapping[str, object], + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict[str, object], str | None]: # mutable-ok: native Messages handler mutates validated headers + headers_out, base = super().validate_anthropic_messages_environment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # validate legacy headers before returning them + headers=go_headers(headers, litellm_params, api_key, model), + model=model, + messages=list(messages), + optional_params=dict(optional_params), + litellm_params=dict(litellm_params), + api_base=go_api_base(api_base, model), + ) + return TypeAdapter(dict[str, object]).validate_python(headers_out), base + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + return ensure_api_route(go_api_base(api_base, model).rstrip("/").removesuffix("/v1"), "v1/messages") diff --git a/litellm/llms/opencode/responses/__init__.py b/litellm/llms/opencode/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/opencode/responses/transformation.py b/litellm/llms/opencode/responses/transformation.py new file mode 100644 index 00000000000..615ca60b45b --- /dev/null +++ b/litellm/llms/opencode/responses/transformation.py @@ -0,0 +1,65 @@ +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.llms.opencode.common_utils import attach_go_reported_cost, go_api_base, go_headers +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + + +class OpenCodeGoResponsesConfig(OpenAIResponsesAPIConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.OPENCODE_GO + + def supports_native_websocket(self) -> bool: + return False + + def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + return ensure_api_route(go_api_base(api_base).rstrip("/").removesuffix("/v1"), "v1/responses") + + def validate_environment( + self, + headers: Mapping[str, object], + model: str, + litellm_params: GenericLiteLLMParams | None, + ) -> dict[str, object]: # mutable-ok: protocol handlers mutate the returned request and headers + return go_headers(headers, litellm_params, litellm_params.api_key if litellm_params else None, model) + + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "Logging", + ) -> ResponsesAPIResponse: + response: Final = super().transform_response_api_response( + model=model, raw_response=raw_response, logging_obj=logging_obj + ) + attach_go_reported_cost(response, response.usage) + return response + + def transform_streaming_response( + self, + model: str, + parsed_chunk: Mapping[str, object], + logging_obj: "Logging", + ) -> ResponsesAPIStreamingResponse: + event: Final = super().transform_streaming_response( # pyright: ignore[reportUnknownMemberType] # base takes an untyped dict + model=model, parsed_chunk=dict(parsed_chunk), logging_obj=logging_obj + ) + if isinstance(event, ResponseCompletedEvent | ResponseFailedEvent | ResponseIncompleteEvent): + attach_go_reported_cost(event.response, event.response.usage) + return event diff --git a/litellm/llms/opencode/transformation.py b/litellm/llms/opencode/transformation.py new file mode 100644 index 00000000000..9b62e7aca95 --- /dev/null +++ b/litellm/llms/opencode/transformation.py @@ -0,0 +1,6 @@ +from litellm.llms.opencode.common_utils import is_opencode_api_base + + +class OpenCodeProviderAdapter: + def get_provider_affinity_header(self, api_base: str | None) -> str | None: + return "x-opencode-session" if is_opencode_api_base(api_base) else None diff --git a/litellm/llms/openrouter/decisions/transformation.py b/litellm/llms/openrouter/decisions/transformation.py index 0fa5c7ca7cb..30d33f09efd 100644 --- a/litellm/llms/openrouter/decisions/transformation.py +++ b/litellm/llms/openrouter/decisions/transformation.py @@ -13,7 +13,7 @@ class OpenRouterDecisionsConfig(BaseDecisionsConfig): return "https://openrouter.ai/api" def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None: - cost: Final = response.usage.extra.get("cost") + cost: Final = response.usage.extra.get("cost") if response.usage is not None else None if not isinstance(cost, (int, float)): return None return float(cost) diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py index 14b0e462ea7..d1530a29515 100644 --- a/litellm/llms/openrouter/embedding/transformation.py +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues from litellm.types.utils import EmbeddingResponse @@ -84,15 +85,7 @@ class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): """ Get the complete URL for OpenRouter Embedding API endpoint. """ - # api_base is already set to https://openrouter.ai/api/v1 in main.py - # Remove trailing slashes - if api_base: - api_base = api_base.rstrip("/") - else: - api_base = "https://openrouter.ai/api/v1" - - # Return the embeddings endpoint - return f"{api_base}/embeddings" + return join_api_base(api_base or "https://openrouter.ai/api/v1", "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index a73a2fe4192..c2d8398aead 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -47,11 +47,12 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles -from pydantic import ConfigDict, RootModel, with_config +from pydantic import ConfigDict, RootModel, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.openrouter.common_utils import OpenRouterException @@ -73,6 +74,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _ImageUrl(TypedDict): @@ -187,11 +190,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): api_base: str | None, litellm_params: dict, ) -> str: - base_url = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" - base_url = base_url.rstrip("/") - if not base_url.endswith("/chat/completions"): - return f"{base_url}/chat/completions" - return base_url + base_url: Final = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" + return ensure_api_route(base_url, "chat/completions") def transform_image_edit_request( self, @@ -250,7 +250,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): logging_obj: LiteLLMLoggingObj, ) -> ImageResponse: try: - response_json: Final[object] = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise OpenRouterException( message=f"Error parsing OpenRouter response: {e}", diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 355df81b87a..a9e07555d6a 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -35,6 +35,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -58,6 +59,7 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) @@ -269,10 +271,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): Default: https://openrouter.ai/api/v1/chat/completions """ if api_base: - if not api_base.endswith("/chat/completions"): - api_base = api_base.rstrip("/") - return f"{api_base}/chat/completions" - return api_base + return ensure_api_route(api_base, "chat/completions") return "https://openrouter.ai/api/v1/chat/completions" @@ -361,7 +360,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ImageResponse: Populated image response """ try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise OpenRouterException( message=f"Error parsing OpenRouter response: {e}", diff --git a/litellm/llms/openrouter/responses/transformation.py b/litellm/llms/openrouter/responses/transformation.py index aeb5b09e14a..e166ccaad2d 100644 --- a/litellm/llms/openrouter/responses/transformation.py +++ b/litellm/llms/openrouter/responses/transformation.py @@ -18,6 +18,7 @@ import litellm from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -138,9 +139,7 @@ class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base or litellm.api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" ) - api_base = api_base.rstrip("/") - - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def supports_native_websocket(self) -> bool: """OpenRouter does not support native WebSocket for Responses API""" diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py index 13c9d805638..d30a9ff534e 100644 --- a/litellm/llms/opensandbox/sandbox/transformation.py +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -20,7 +20,9 @@ from litellm.constants import ( OPEN_SANDBOX_POLL_INTERVAL, OPEN_SANDBOX_READY_TIMEOUT, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.sandbox.transformation import ( SANDBOX_MAX_OUTPUT_BYTES, BaseSandboxConfig, @@ -43,6 +45,13 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _JSON_VALUE: Final = TypeAdapter(object) +def _decoded_json_object(value: object) -> Mapping[str, object]: + mapping: Final = as_str_mapping(value) + if mapping is not None: + return mapping + raise AttributeError(f"{type(value).__name__!r} object has no attribute 'get'", name="get", obj=value) + + class OpenSandboxSandboxConfig(BaseSandboxConfig): def _http(self, client: AsyncHTTPHandler | None) -> AsyncHTTPHandler: if client is not None: @@ -93,7 +102,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): ) response: Final = await self._http(client).post( - url=f"{base}/sandboxes", + url=join_api_base(base, "sandboxes"), headers=self._lifecycle_headers(key), json=body, ) @@ -192,7 +201,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): key: Final = self._api_key(api_key=api_key, handle=handle) try: response: Final = await self._http(client).delete( - url=f"{base}/sandboxes/{handle.id}", + url=join_api_base(base, f"sandboxes/{handle.id}"), headers=self._lifecycle_headers(key), ) except httpx.HTTPStatusError as e: @@ -258,7 +267,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): deadline: Final = time.monotonic() + ready_timeout while True: response = await self._http(client).get( - url=f"{api_base}/sandboxes/{sandbox_id}", + url=join_api_base(api_base, f"sandboxes/{sandbox_id}"), headers=headers, ) data = _JSON_VALUE.validate_python(response.json()) @@ -316,11 +325,11 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): client: AsyncHTTPHandler | None, ) -> tuple[str, dict[str, str]]: response: Final = await self._http(client).get( - url=f"{api_base}/sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}", + url=join_api_base(api_base, f"sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}"), headers=headers, params={"use_server_proxy": use_server_proxy}, ) - data: Final = response.json() + data: Final = _decoded_json_object(_JSON_VALUE.validate_python(response.json())) endpoint: Final = data.get("endpoint") if not endpoint: raise ValueError(f"OpenSandbox did not return an execd endpoint for {sandbox_id}") @@ -491,7 +500,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): if not data: return None try: - parsed: Final = json.loads(data) + parsed: Final = _JSON_VALUE.validate_python(json.loads(data)) except json.JSONDecodeError: return None if not isinstance(parsed, dict): diff --git a/litellm/llms/ovhcloud/audio_transcription/transformation.py b/litellm/llms/ovhcloud/audio_transcription/transformation.py index 3ae5f6fa69c..dd61bfb0c13 100644 --- a/litellm/llms/ovhcloud/audio_transcription/transformation.py +++ b/litellm/llms/ovhcloud/audio_transcription/transformation.py @@ -12,6 +12,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -63,9 +64,8 @@ class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") - complete_url: Final = f"{api_base}/audio/transcriptions" - return complete_url + base_url: Final = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base + return join_api_base(base_url, "audio/transcriptions") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return OVHCloudException( diff --git a/litellm/llms/ovhcloud/chat/transformation.py b/litellm/llms/ovhcloud/chat/transformation.py index 3d3953eed34..91b3dbe6268 100644 --- a/litellm/llms/ovhcloud/chat/transformation.py +++ b/litellm/llms/ovhcloud/chat/transformation.py @@ -9,6 +9,7 @@ from typing import Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig @@ -31,9 +32,8 @@ class OVHCloudChatConfig(OpenAIGPTConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") - complete_url: Final = f"{api_base}/chat/completions" - return complete_url + base_url: Final = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base + return join_api_base(base_url, "chat/completions") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return OVHCloudException( diff --git a/litellm/llms/ovhcloud/embedding/transformation.py b/litellm/llms/ovhcloud/embedding/transformation.py index cdf095e6ad1..487c40afa52 100644 --- a/litellm/llms/ovhcloud/embedding/transformation.py +++ b/litellm/llms/ovhcloud/embedding/transformation.py @@ -10,6 +10,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -35,9 +36,8 @@ class OVHCloudEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") - complete_url: Final = f"{api_base}/embeddings" - return complete_url + base_url: Final = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base + return join_api_base(base_url, "embeddings") def validate_environment( self, diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 885fc98ffea..dea8bcb73a4 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -6,11 +6,14 @@ It uses the field targeting configuration from litellm_logging_obj to extract specific fields for guardrail processing. """ -from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional, Protocol, runtime_checkable +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import without_server_streaming_classification +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -87,7 +90,7 @@ class PassThroughEndpointHandler(BaseTranslation): payload_to_check: Final = { k: v for k, v in without_server_streaming_classification(data).items() - if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj") + if not k.startswith("_") and k not in ("metadata", "litellm_metadata", "litellm_logging_obj") } verbose_proxy_logger.debug("PassThroughEndpointHandler: Using full payload for guardrail") return safe_dumps(payload_to_check) @@ -215,6 +218,19 @@ class _EventStreamEndpointChecker(Protocol): def event_stream_endpoint_is_de_anonymizable(self, endpoint: str, model: str, body: object) -> bool: ... +class _EventStreamDeAnonymizer(Protocol): + async def de_anonymize_event_stream( + self, + body_bytes: bytes, + proxy_logging_obj: "ProxyLogging", + user_api_key_dict: "UserAPIKeyAuth", + data: Mapping[str, object], + ) -> bytes: ... + + +_TRANSLATION_RESULT: Final = TypeAdapter(object) + + _PROVIDER_HANDLERS: dict[str, type[BaseTranslation]] = {} @@ -260,10 +276,12 @@ class LlmPassthroughRouteHandler(BaseTranslation): provider, ) return data - return await handler_cls().process_input_messages( - data=data, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, + return _TRANSLATION_RESULT.validate_python( + await handler_cls().process_input_messages( + data=data, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + ) ) async def process_output_response( @@ -282,12 +300,14 @@ class LlmPassthroughRouteHandler(BaseTranslation): provider, ) return response - return await handler_cls().process_output_response( - response=response, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=request_data, + return _TRANSLATION_RESULT.validate_python( + await handler_cls().process_output_response( + response=response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) ) @staticmethod @@ -399,9 +419,11 @@ class LlmPassthroughRouteHandler(BaseTranslation): return handler_cls.event_stream_media_type() @staticmethod - def _resolve_event_stream_de_anonymizer(provider: str | None): + def _resolve_event_stream_de_anonymizer(provider: str | None) -> Callable[..., Awaitable[bytes]] | None: handler_cls: Final = _get_provider_handlers().get(provider or "") - return getattr(handler_cls, "de_anonymize_event_stream", None) + if not has_attributes(handler_cls, _EventStreamDeAnonymizer, "de_anonymize_event_stream"): + return None + return handler_cls.de_anonymize_event_stream @staticmethod def supports_event_stream_relay(provider: object, endpoint: object, model: object, body: object) -> bool: diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py index 945a58d4d75..7b4da933976 100644 --- a/litellm/llms/perplexity/embedding/transformation.py +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -20,6 +20,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -68,7 +69,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): if api_base: clean_api_base: Final = api_base.rstrip("/") if not clean_api_base.endswith("/embeddings"): - return f"{clean_api_base}/v1/embeddings" + return join_api_base(clean_api_base, "v1/embeddings") return clean_api_base return "https://api.perplexity.ai/v1/embeddings" diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index 91d40d8247f..77817a5e5bb 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -12,8 +12,10 @@ Ref: https://docs.perplexity.ai/api-reference/responses-post from typing import Any, Final import httpx +from pydantic import TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str @@ -21,6 +23,8 @@ from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +_DECODED_JSON: Final = TypeAdapter(object) + class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): def get_supported_openai_params(self, model: str) -> list: @@ -51,7 +55,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str: api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" - return f"{api_base.rstrip('/')}/v1/responses" + return join_api_base(api_base, "v1/responses") def _ensure_message_type(self, input: str | ResponseInputParam) -> str | ResponseInputParam: """Ensure list input items have type='message' (required by Perplexity).""" @@ -103,7 +107,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): ) -> ResponsesAPIResponse: """Check for Perplexity's status:'failed' on HTTP 200 before delegating to base.""" try: - raw_response_json = raw_response.json() + raw_response_json = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raw_response_json = None diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index b46e4225309..1f3b0fe729d 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -9,6 +9,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -78,11 +79,7 @@ class PerplexitySearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("PERPLEXITY_API_BASE") or self.PERPLEXITY_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("PERPLEXITY_API_BASE") or self.PERPLEXITY_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/petals/completion/transformation.py b/litellm/llms/petals/completion/transformation.py index ee20c2b12d5..a875992448a 100644 --- a/litellm/llms/petals/completion/transformation.py +++ b/litellm/llms/petals/completion/transformation.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final from httpx import Headers, Response @@ -57,7 +59,7 @@ class PetalsConfig(BaseConfig): top_p: float | None = None, repetition_penalty: float | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index e4b06c36bf4..89851f60d76 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -1,6 +1,6 @@ from typing import TYPE_CHECKING, Any, Final -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -66,10 +66,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): "PG Vector API base URL is required. Set PG_VECTOR_API_BASE environment variable or pass api_base in litellm_params." ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/v1/vector_stores" + return join_api_base(api_base, "v1/vector_stores") def transform_search_vector_store_request( self, @@ -79,10 +76,10 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}/search" + url: Final = join_api_base(api_base, f"{encoded_vector_store_id}/search") _, request_body = super().transform_search_vector_store_request( vector_store_id=vector_store_id, query=query, diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 0dd7e27722e..b30713fa119 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( prompt_factory, render_custom_prompt_template, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -337,12 +338,12 @@ class PredibaseConfig(BaseConfig): elif "PREDIBASE_API_BASE" in os.environ: base_url = os.getenv("PREDIBASE_API_BASE", "") - completion_url = f"{base_url}/{tenant_id}/deployments/v2/llms/{model}" + completion_url = join_api_base(base_url, f"{tenant_id}/deployments/v2/llms/{model}") should_stream: Final = stream if stream is not None else optional_params.get("stream", False) if should_stream is True: - completion_url += "/generate_stream" + completion_url = join_api_base(completion_url, "generate_stream") else: - completion_url += "/generate" + completion_url = join_api_base(completion_url, "generate") return completion_url def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: diff --git a/litellm/llms/ragflow/chat/transformation.py b/litellm/llms/ragflow/chat/transformation.py index 97b723f7dd5..f626cd8e55f 100644 --- a/litellm/llms/ragflow/chat/transformation.py +++ b/litellm/llms/ragflow/chat/transformation.py @@ -13,7 +13,7 @@ Model name format: from typing import Final import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.openai.openai import OpenAIConfig from litellm.secret_managers.main import get_secret, get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -122,7 +122,7 @@ class RAGFlowConfig(OpenAIConfig): if not path.startswith("/"): path = "/" + path - return f"{api_base}{path}" + return join_api_base(api_base, path) def _get_openai_compatible_provider_info( self, diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index 007b047cc1b..ef44b6e4c0f 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, RootModel, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -95,15 +96,16 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): - api_base in litellm_params - Default: http://localhost:9380 """ - api_base = ( - api_base or litellm_params.get("api_base") or get_secret_str("RAGFLOW_API_BASE") or "http://localhost:9380" + params_api_base: Final = litellm_params.get("api_base") + resolved_api_base: Final = ( + api_base + or (f"{params_api_base}" if params_api_base else None) + or get_secret_str("RAGFLOW_API_BASE") + or "http://localhost:9380" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - # RAGFlow datasets API endpoint - return f"{api_base}/api/v1/datasets" + return join_api_base(resolved_api_base, "api/v1/datasets") def transform_search_vector_store_request( self, diff --git a/litellm/llms/recraft/image_edit/transformation.py b/litellm/llms/recraft/image_edit/transformation.py index 46905580382..74bfa52c5e0 100644 --- a/litellm/llms/recraft/image_edit/transformation.py +++ b/litellm/llms/recraft/image_edit/transformation.py @@ -5,6 +5,7 @@ import httpx from httpx._types import RequestFiles from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -68,11 +69,8 @@ class RecraftImageEditConfig(BaseImageEditConfig): Some providers need `model` in `api_base` """ - complete_url: str = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - complete_url = f"{complete_url}/{self.IMAGE_EDIT_ENDPOINT}" - return complete_url + base_url: Final = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, self.IMAGE_EDIT_ENDPOINT) def validate_environment( self, diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py index 0fcfafb79fb..49a2938e1c0 100644 --- a/litellm/llms/recraft/image_generation/transformation.py +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -25,6 +26,7 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class RecraftImageGenerationConfig(BaseImageGenerationConfig): @@ -72,11 +74,8 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" - return complete_url + base_url: Final = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, self.IMAGE_GENERATION_ENDPOINT) def validate_environment( self, @@ -136,7 +135,7 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): https://www.recraft.ai/docs#generate-image """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index ddc1b9474ac..23c8c86e7ca 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -1,9 +1,10 @@ import asyncio import time from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, cast import httpx +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger @@ -11,6 +12,8 @@ from litellm.constants import ( RUNWAYML_DEFAULT_API_VERSION, RUNWAYML_POLLING_TIMEOUT, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -35,9 +38,25 @@ class _RunwayMLTask(TypedDict, total=False): id: ReadOnly[str] status: ReadOnly[str] - output: ReadOnly[Sequence[str | Mapping[str, str]]] - failure: ReadOnly[str] - failureCode: ReadOnly[str] + output: ReadOnly[Sequence[str | Mapping[str, str | None]] | None] + failure: ReadOnly[str | None] + failureCode: ReadOnly[str | None] + + +_JSON_VALUE: Final = TypeAdapter(object) + + +def _decoded_json(response: httpx.Response) -> object: + return _JSON_VALUE.validate_python(response.json()) + + +def _runway_task(decoded: object) -> _RunwayMLTask: + task: Final = decoded_json_object(decoded) + return cast(_RunwayMLTask, task) # cast-ok: the payload passes through unvalidated, as before it was typed + + +def _parse_runway_task(response: httpx.Response) -> _RunwayMLTask: + return _runway_task(_decoded_json(response)) class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): @@ -62,12 +81,8 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - if self.IMAGE_GENERATION_ENDPOINT: - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" - return complete_url + base_url: Final = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, self.IMAGE_GENERATION_ENDPOINT) def validate_environment( self, @@ -226,8 +241,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): start_time: Final = time.time() # Build task status URL - api_base = api_base.rstrip("/") - task_url: Final = f"{api_base}/v1/tasks/{task_id}" + task_url: Final = join_api_base(api_base, f"v1/tasks/{task_id}") verbose_logger.debug("Polling RunwayML task: %s", task_url) @@ -238,7 +252,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): response = client.get(url=task_url, headers=headers) response.raise_for_status() - response_data: _RunwayMLTask = response.json() + response_data: _RunwayMLTask = _parse_runway_task(response) # Check task status status = self._check_task_status(response_data=response_data) @@ -275,8 +289,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): start_time: Final = time.time() # Build task status URL - api_base = api_base.rstrip("/") - task_url: Final = f"{api_base}/v1/tasks/{task_id}" + task_url: Final = join_api_base(api_base, f"v1/tasks/{task_id}") verbose_logger.debug("Polling RunwayML task (async): %s", task_url) @@ -287,7 +300,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): response = await client.get(url=task_url, headers=headers) response.raise_for_status() - response_data: _RunwayMLTask = response.json() + response_data: _RunwayMLTask = _parse_runway_task(response) # Check task status status = self._check_task_status(response_data=response_data) @@ -333,13 +346,14 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): } """ try: - response_data: _RunwayMLTask = raw_response.json() + decoded: Final = _decoded_json(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) + response_data: _RunwayMLTask = _runway_task(decoded) verbose_logger.debug("RunwayML starting polling...") @@ -363,7 +377,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) # Update response_data with polled result - response_data = raw_response.json() + response_data = _parse_runway_task(raw_response) verbose_logger.debug("RunwayML polling complete, transforming to OpenAI format") @@ -393,13 +407,14 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): We need to poll the task until it completes (status SUCCEEDED) using async polling. """ try: - response_data: _RunwayMLTask = raw_response.json() + decoded: Final = _decoded_json(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) + response_data: _RunwayMLTask = _runway_task(decoded) verbose_logger.debug("RunwayML starting polling (async)...") @@ -423,7 +438,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) # Update response_data with polled result - response_data = raw_response.json() + response_data = _parse_runway_task(raw_response) verbose_logger.debug("RunwayML polling complete (async), transforming to OpenAI format") diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index d41ffea17dc..ce1ec1714f1 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -18,6 +18,7 @@ from litellm.constants import ( RUNWAYML_DEFAULT_API_VERSION, RUNWAYML_POLLING_TIMEOUT, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -40,6 +41,11 @@ class _RunwayTtsTaskResponse(TypedDict, total=False): failureCode: ReadOnly[str] +def _parse_runway_tts_task(response: httpx.Response) -> _RunwayTtsTaskResponse: + task: Final[_RunwayTtsTaskResponse] = response.json() + return task + + class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ Configuration for RunwayML Text-to-Speech @@ -230,10 +236,8 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ Get the complete URL for RunwayML TTS request """ - complete_url = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - return f"{complete_url}/{self.TTS_ENDPOINT_PATH}" + base_url: Final = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, self.TTS_ENDPOINT_PATH) @staticmethod def _check_timeout(start_time: float, timeout_secs: float) -> None: @@ -311,8 +315,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): start_time: Final = time.time() # Build task status URL - api_base = api_base.rstrip("/") - task_url: Final = f"{api_base}/v1/tasks/{task_id}" + task_url: Final = join_api_base(api_base, f"v1/tasks/{task_id}") verbose_logger.debug("Polling RunwayML TTS task: %s", task_url) @@ -323,7 +326,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): response = client.get(url=task_url, headers=headers) response.raise_for_status() - response_data: _RunwayTtsTaskResponse = response.json() + response_data: _RunwayTtsTaskResponse = _parse_runway_tts_task(response) # Check task status status = self._check_task_status(response_data=response_data) @@ -359,8 +362,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): start_time: Final = time.time() # Build task status URL - api_base = api_base.rstrip("/") - task_url: Final = f"{api_base}/v1/tasks/{task_id}" + task_url: Final = join_api_base(api_base, f"v1/tasks/{task_id}") verbose_logger.debug("Polling RunwayML TTS task (async): %s", task_url) @@ -371,7 +373,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): response = await client.get(url=task_url, headers=headers) response.raise_for_status() - response_data: _RunwayTtsTaskResponse = response.json() + response_data: _RunwayTtsTaskResponse = _parse_runway_tts_task(response) # Check task status status = self._check_task_status(response_data=response_data) @@ -462,7 +464,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_data: Final[_RunwayTtsTaskResponse] = raw_response.json() + response_data: Final[_RunwayTtsTaskResponse] = _parse_runway_tts_task(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error parsing RunwayML TTS response: {e}", @@ -492,7 +494,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get the completed task data - task_data: Final[_RunwayTtsTaskResponse] = polled_response.json() + task_data: Final[_RunwayTtsTaskResponse] = _parse_runway_tts_task(polled_response) verbose_logger.debug("RunwayML TTS polling complete, downloading audio") @@ -531,7 +533,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_data: Final[_RunwayTtsTaskResponse] = raw_response.json() + response_data: Final[_RunwayTtsTaskResponse] = _parse_runway_tts_task(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error parsing RunwayML TTS response: {e}", @@ -561,7 +563,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get the completed task data - task_data: Final[_RunwayTtsTaskResponse] = polled_response.json() + task_data: Final[_RunwayTtsTaskResponse] = _parse_runway_tts_task(polled_response) verbose_logger.debug("RunwayML TTS polling complete (async), downloading audio") diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 8b765b1390b..ad8a24c1535 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -9,7 +9,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( @@ -119,7 +119,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): @staticmethod def _parse_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse: - return raw_response.json() + return _parse_runway_task_response(raw_response) def get_supported_openai_params(self, model: str) -> list: """ @@ -274,7 +274,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): files_list: Final[RequestFiles] = [] - return request_data, files_list, f"{api_base}/{endpoint}" + return request_data, files_list, join_api_base(api_base, endpoint) def _select_generation_endpoint(self, request_data: Mapping[str, object]) -> str: if request_data.get("promptVideo") is not None or request_data.get("videoUri") is not None: @@ -413,7 +413,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Get task status to retrieve video URL - url: Final = f"{api_base}/tasks/{encoded_video_id}" + url: Final = join_api_base(api_base, f"tasks/{encoded_video_id}") return url, dict[str, str]() @@ -571,7 +571,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for task cancellation - url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel" + url: Final = join_api_base(api_base, f"tasks/{encoded_video_id}/cancel") return url, dict[str, str]() @@ -608,7 +608,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the full URL for task status retrieval - url: Final = f"{api_base}/tasks/{encoded_video_id}" + url: Final = join_api_base(api_base, f"tasks/{encoded_video_id}") # Empty dict for GET request (no body) return url, dict[str, str]() diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 044315168d2..2b0de3dc9ea 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, @@ -131,7 +132,7 @@ class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM litellm_logging_obj: LiteLLMLoggingObj, ) -> tuple[str, dict[str, object]]: litellm_logging_obj.model_call_details["query"] = query_text - return f"{api_base}/QueryVectors", { + return join_api_base(api_base, "QueryVectors"), { "vectorBucketName": bucket_name, "indexName": index_name, "queryVector": {"float32": list(query_vector)}, diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 14f0e3ead62..eaa805e7769 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -1,3 +1,4 @@ +import builtins import codecs import functools import json @@ -34,6 +35,14 @@ class _SageMakerMessagesChunkData(TypedDict, total=False): choices: ReadOnly[Required[list[Mapping[str, object]]]] +class _SageMakerPayloadPart(TypedDict, total=False): + bytes: ReadOnly[builtins.bytes] + + +class _SageMakerStreamEvent(TypedDict, total=False): + chunk: ReadOnly[_SageMakerPayloadPart] + + class _SageMakerChunkDataRoot(RootModel[_SageMakerChunkData]): model_config = ConfigDict(frozen=True, hide_input_in_errors=True) @@ -345,7 +354,7 @@ class AWSEventStreamDecoder: ), ) response_dict: Final = event.to_response_dict() - parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) + parsed_response: Final[_SageMakerStreamEvent] = self.parser.parse(response_dict, response_stream_shape) body: Final = response_dict.get("body") if response_dict["status_code"] != 200: diff --git a/litellm/llms/sagemaker/embedding/cohere_transformation.py b/litellm/llms/sagemaker/embedding/cohere_transformation.py index 35352f65dda..8c0f4437edf 100644 --- a/litellm/llms/sagemaker/embedding/cohere_transformation.py +++ b/litellm/llms/sagemaker/embedding/cohere_transformation.py @@ -10,7 +10,8 @@ be of type Object`. Reference: https://docs.cohere.com/v2/reference/embed """ -from typing import TYPE_CHECKING, Any, cast +from collections.abc import Sequence +from typing import TYPE_CHECKING, cast if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -25,6 +26,7 @@ from litellm.llms.bedrock.embed.cohere_transformation import ( BedrockCohereEmbeddingConfig, ) from litellm.llms.cohere.embed.v1_transformation import CohereEmbeddingConfig +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import EmbeddingResponse from ..common_utils import SagemakerError @@ -123,7 +125,7 @@ class SagemakerCohereEmbeddingConfig(BaseEmbeddingConfig): self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/sagemaker/embedding/openai_transformation.py b/litellm/llms/sagemaker/embedding/openai_transformation.py index a3f0ebfbf59..1222c697832 100644 --- a/litellm/llms/sagemaker/embedding/openai_transformation.py +++ b/litellm/llms/sagemaker/embedding/openai_transformation.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllEmbeddingInputValues from httpx._models import Headers, Response -from pydantic import BaseModel, ConfigDict, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig @@ -25,6 +25,7 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import STR_KEYED_MAPPING, SagemakerError SUPPORTED_OPENAI_PARAMS: Final = ("dimensions", "encoding_format", "user") +_DECODED_JSON: Final = TypeAdapter(object) class OpenAIEmbeddingUsage(BaseModel): @@ -94,7 +95,7 @@ class SagemakerOpenAIEmbeddingConfig(BaseEmbeddingConfig): litellm_params: Mapping[str, object] | None = None, ) -> "EmbeddingResponse": try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except ValueError as e: raise SagemakerError( message=f"Failed to parse response: {e}", diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 1395f9b98ec..2cd02683167 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -4,7 +4,8 @@ Translate from OpenAI's `/v1/embeddings` to Sagemaker's `/invoke` In the Huggingface TGI format. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Sequence +from typing import TYPE_CHECKING, Final if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -15,6 +16,7 @@ from httpx._models import Headers, Response from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import SagemakerError @@ -158,7 +160,7 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/sambanova/embedding/transformation.py b/litellm/llms/sambanova/embedding/transformation.py index cf39c0a092f..93cbb13ac47 100644 --- a/litellm/llms/sambanova/embedding/transformation.py +++ b/litellm/llms/sambanova/embedding/transformation.py @@ -10,6 +10,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -37,11 +38,7 @@ class SambaNovaEmbeddingConfig(BaseEmbeddingConfig): ) -> str: if api_base is None: raise ValueError("api_base is required for SambaNova embeddings") - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/embeddings"): - api_base = f"{api_base}/embeddings" - return api_base + return ensure_api_route(api_base, "embeddings") def validate_environment( self, diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index a82dfe72322..622c28f3c90 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import BaseModel import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -24,7 +25,7 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -from ..credentials import get_token_creator +from ..credentials import AICoreConfiguration, AICoreDeploymentList, get_token_creator from .handler import ( AsyncSAPStreamIterator, GenAIHubOrchestrationError, @@ -185,12 +186,14 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): # Keep a short, tight client lifecycle here to avoid fd leaks client: Final = litellm.module_level_client # with httpx.Client(timeout=30) as client: - deployments: Final = client.get(f"{self.base_url}/lm/deployments", headers=self.headers).json() + deployments: Final[AICoreDeploymentList] = client.get( + join_api_base(self.base_url, "lm/deployments"), headers=self.headers + ).json() valid: Final[list[tuple[str, str]]] = [] for dep in deployments.get("resources", []): if dep.get("scenarioId") == "orchestration": - cfg = client.get( - f"{self.base_url}/lm/configurations/{dep['configurationId']}", + cfg: AICoreConfiguration = client.get( + join_api_base(self.base_url, f"lm/configurations/{dep['configurationId']}"), headers=self.headers, ).json() if cfg.get("executableId") == "orchestration": @@ -263,7 +266,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): litellm_params: dict, stream: bool | None = None, ): - api_base_: Final = f"{self.deployment_url}/v2/completion" + api_base_: Final = join_api_base(self.deployment_url, "v2/completion") return api_base_ def _build_prompt_module( diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index b1618f38623..74d065956cc 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path from threading import Lock -from typing import Any, Final, Protocol +from typing import Final, Protocol import httpx from pydantic import TypeAdapter @@ -16,6 +16,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_httpx_client AUTH_ENDPOINT_SUFFIX: Final = "/oauth/token" @@ -33,6 +34,21 @@ DEFAULT_HOME_PATH: Final = os.path.join(os.path.expanduser("~"), ".aicore") _DECODED_JSON: Final = TypeAdapter(object) +class AICoreDeployment(TypedDict): + configurationId: ReadOnly[str] + scenarioId: ReadOnly[str] + deploymentUrl: ReadOnly[str] + createdAt: ReadOnly[str] + + +class AICoreDeploymentList(TypedDict, total=False): + resources: ReadOnly[Sequence[AICoreDeployment]] + + +class AICoreConfiguration(TypedDict): + executableId: ReadOnly[str] + + def _get_home() -> str: return os.getenv(HOME_PATH_ENV_VAR, DEFAULT_HOME_PATH) @@ -58,7 +74,7 @@ def _get_nested(d: object, path: Sequence[str]) -> object: return cur -def _load_json_env(var_name: str) -> dict[str, object] | None: +def _load_json_env(var_name: str) -> Mapping[str, Sequence[Mapping[str, object]]] | None: raw: Final = os.environ.get(var_name) if not raw: return None @@ -75,7 +91,7 @@ def _str_or_none(value: object) -> str | None: return None -def _load_vcap() -> dict[str, Any]: +def _load_vcap() -> Mapping[str, Sequence[Mapping[str, object]]]: return _load_json_env(VCAP_SERVICES_ENV_VAR) or {} @@ -112,7 +128,7 @@ CREDENTIAL_VALUES: Final[list[CredentialsValue]] = [ CredentialsValue( "base_url", ("serviceurls", "AI_API_URL"), - transform_fn=lambda url: url.rstrip("/") + ("" if url.endswith("/v2") else "/v2"), + transform_fn=lambda url: ensure_api_route(url, "v2"), ), CredentialsValue( "cert_url", @@ -195,7 +211,7 @@ def resolve_resource_group(sources: list[Source]) -> str | None: def _parse_service_key_once( service_key: str | dict | None, -) -> dict[str, object] | None: +) -> object: """ Pre-parse service_key if it's a string to avoid repeated JSON parsing. @@ -207,7 +223,7 @@ def _parse_service_key_once( return service_key if isinstance(service_key, str): try: - return json.loads(service_key) + return _DECODED_JSON.validate_python(json.loads(service_key)) except json.JSONDecodeError: verbose_logger.warning("SAP service key is a string but not valid JSON. Skipping this source.") return None @@ -217,7 +233,7 @@ def _parse_service_key_once( return None -def _resolve_credential_from_service_key(service_key: str | dict | None, cv: CredentialsValue) -> str | None: +def _resolve_credential_from_service_key(service_key: object, cv: CredentialsValue) -> str | None: if service_key is None: return None val: Final = _str_or_none(_get_nested(service_key, (("credentials",) + cv.vcap_key) if cv.vcap_key else (cv.name,))) @@ -252,14 +268,16 @@ def fetch_credentials( """ config: Final = init_conf(profile) - service_key = _parse_service_key_once(service_key or litellm.sap_service_key or os.environ.get(SERVICE_KEY_ENV_VAR)) + parsed_service_key: Final = _parse_service_key_once( + service_key or litellm.sap_service_key or os.environ.get(SERVICE_KEY_ENV_VAR) + ) vcap_service: Final = _get_vcap_service(VCAP_AICORE_SERVICE_NAME) sources: Final = [ Source("kwargs", lambda cv: _str_or_none(kwargs.get(cv.name))), Source( "service key", - lambda cv: _resolve_credential_from_service_key(service_key, cv), + lambda cv: _resolve_credential_from_service_key(parsed_service_key, cv), ), Source( "environment variables", diff --git a/litellm/llms/sap/embed/transformation.py b/litellm/llms/sap/embed/transformation.py index e91dfb91081..a59b7e1e00d 100644 --- a/litellm/llms/sap/embed/transformation.py +++ b/litellm/llms/sap/embed/transformation.py @@ -8,6 +8,7 @@ from typing import Final, Literal import httpx from pydantic import Field +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.embedding.transformation import ( BaseEmbeddingConfig, LiteLLMLoggingObj, @@ -18,7 +19,7 @@ from litellm.types.llms.openai import AllEmbeddingInputValues from litellm.types.utils import EmbeddingResponse from ..chat.handler import GenAIHubOrchestrationError -from ..credentials import get_token_creator +from ..credentials import AICoreConfiguration, AICoreDeploymentList, get_token_creator class Usage(LiteLLMBaseModel): @@ -101,11 +102,13 @@ class GenAIHubEmbeddingConfig(BaseEmbeddingConfig): def deployment_url(self) -> str: with httpx.Client(timeout=30) as client: valid_deployments: Final = [] - deployments: Final = client.get(self.base_url + "/lm/deployments", headers=self.headers).json() + deployments: Final[AICoreDeploymentList] = client.get( + join_api_base(self.base_url, "lm/deployments"), headers=self.headers + ).json() for deployment in deployments.get("resources", []): if deployment["scenarioId"] == "orchestration": - config_details = client.get( - self.base_url + f"/lm/configurations/{deployment['configurationId']}", + config_details: AICoreConfiguration = client.get( + join_api_base(self.base_url, f"lm/configurations/{deployment['configurationId']}"), headers=self.headers, ).json() if config_details["executableId"] == "orchestration": @@ -144,7 +147,7 @@ class GenAIHubEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - url: Final = self.deployment_url.rstrip("/") + "/v2/embeddings" + url: Final = join_api_base(self.deployment_url, "v2/embeddings") return url def transform_embedding_request( diff --git a/litellm/llms/scaledown/chat/transformation.py b/litellm/llms/scaledown/chat/transformation.py index 44eb4785615..3224bc97148 100644 --- a/litellm/llms/scaledown/chat/transformation.py +++ b/litellm/llms/scaledown/chat/transformation.py @@ -9,6 +9,7 @@ from urllib.parse import unquote import httpx from pydantic import JsonValue, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -103,8 +104,8 @@ class ScaleDownChatConfig(BaseConfig): base: Final = _root(api_base or _trusted_root()) operation: Final = _operation(model) if operation in DECISIONS_MODELS: - return f"{base}/v1/scaledown" - return f"{base}{NATIVE_PATHS[operation]}" + return join_api_base(base, "v1/scaledown") + return join_api_base(base, NATIVE_PATHS[operation]) def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig signature operation: Final = _operation(model) diff --git a/litellm/llms/scaleway/audio_transcription/transformation.py b/litellm/llms/scaleway/audio_transcription/transformation.py index 7ebf4475355..8913cc362b3 100644 --- a/litellm/llms/scaleway/audio_transcription/transformation.py +++ b/litellm/llms/scaleway/audio_transcription/transformation.py @@ -9,6 +9,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -58,8 +59,7 @@ class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/") - return f"{api_base}/audio/transcriptions" + return join_api_base("https://api.scaleway.ai/v1" if api_base is None else api_base, "audio/transcriptions") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return ScalewayAudioTranscriptionException( diff --git a/litellm/llms/scaleway/rerank/transformation.py b/litellm/llms/scaleway/rerank/transformation.py index 921252d7091..1de5d71f17a 100644 --- a/litellm/llms/scaleway/rerank/transformation.py +++ b/litellm/llms/scaleway/rerank/transformation.py @@ -9,6 +9,7 @@ API reference: https://www.scaleway.com/en/developers/api/generative-apis/#path- from collections.abc import Mapping from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.jina_ai.rerank.transformation import JinaAIRerankConfig from litellm.secret_managers.main import get_secret_str @@ -25,8 +26,8 @@ class ScalewayRerankConfig(JinaAIRerankConfig): model: str, optional_params: Mapping[str, object] | None = None, ) -> str: - base: Final = SCALEWAY_API_BASE if api_base is None else api_base.rstrip("/") - return f"{base}/rerank" + base: Final = SCALEWAY_API_BASE if api_base is None else api_base + return join_api_base(base, "rerank") def validate_environment( self, diff --git a/litellm/llms/searxng/search/transformation.py b/litellm/llms/searxng/search/transformation.py index 265a3db72f5..b1753821cd2 100644 --- a/litellm/llms/searxng/search/transformation.py +++ b/litellm/llms/searxng/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -101,12 +102,7 @@ class SearXNGSearchConfig(BaseSearchConfig): "or pass `api_base` parameter. Example: os.environ['SEARXNG_API_BASE'] = 'https://your-searxng-instance.com'" ) - # Append "/search" to the api base if it's not already there - if not api_base.endswith("/search"): - if api_base.endswith("/"): - api_base = f"{api_base}search" - else: - api_base = f"{api_base}/search" + api_base = ensure_api_route(api_base, "search") # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_searxng_params" in data: diff --git a/litellm/llms/serper/search/transformation.py b/litellm/llms/serper/search/transformation.py index 2539610ea59..e82427e4d0c 100644 --- a/litellm/llms/serper/search/transformation.py +++ b/litellm/llms/serper/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -84,12 +85,7 @@ class SerperSearchConfig(BaseSearchConfig): Get complete URL for Search endpoint. """ api_base = api_base or get_secret_str("SERPER_API_BASE") or self.SERPER_API_BASE - api_base = api_base.rstrip("/") - - if not api_base.endswith("/search"): - api_base = f"{api_base}/search" - - return api_base + return ensure_api_route(api_base, "search") def transform_search_request( self, diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 6a480904a3e..030413030eb 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( select_anthropic_content_block_type_for_file, ) from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload @@ -239,8 +240,8 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): ) -> str: api_base = self._get_api_base(api_base, optional_params) if is_snowflake_claude_model(model): - return f"{api_base}/cortex/v1/messages" - return f"{api_base}/cortex/v1/chat/completions" + return join_api_base(api_base, "cortex/v1/messages") + return join_api_base(api_base, "cortex/v1/chat/completions") def validate_environment( self, diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py index ec482337e52..45bec1283a3 100644 --- a/litellm/llms/snowflake/embedding/transformation.py +++ b/litellm/llms/snowflake/embedding/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues @@ -31,7 +32,7 @@ class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): ) -> str: api_base = self._get_api_base(api_base, optional_params) - return f"{api_base}/cortex/inference:embed" + return join_api_base(api_base, "cortex/inference:embed") def transform_embedding_request( self, diff --git a/litellm/llms/snowflake/utils.py b/litellm/llms/snowflake/utils.py index a8c14d1d7f3..7cc88ab66e1 100644 --- a/litellm/llms/snowflake/utils.py +++ b/litellm/llms/snowflake/utils.py @@ -1,6 +1,7 @@ import re from typing import TYPE_CHECKING, Any, Final +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -53,7 +54,7 @@ class SnowflakeBaseConfig: optional_params[param] = value return optional_params - def _get_api_base(self, api_base, optional_params): + def _get_api_base(self, api_base: str | None, optional_params) -> str: if not api_base: if "account_id" in optional_params: account_id = optional_params.pop("account_id") @@ -65,10 +66,7 @@ class SnowflakeBaseConfig: raise ValueError("Invalid account_id format") api_base = f"https://{account_id}.snowflakecomputing.com/api/v2" - api_base = api_base.rstrip("/") - if not api_base.endswith("/api/v2"): - api_base += "/api/v2" - return api_base + return ensure_api_route(api_base, "api/v2") def validate_environment( self, diff --git a/litellm/llms/soniox/audio_transcription/handler.py b/litellm/llms/soniox/audio_transcription/handler.py index f86bc142898..a48e3e3ed4c 100644 --- a/litellm/llms/soniox/audio_transcription/handler.py +++ b/litellm/llms/soniox/audio_transcription/handler.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_name, process_audio_file, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -329,7 +330,7 @@ class SonioxAudioTranscriptionHandler: input=None, api_key=api_key, additional_args={ - "api_base": f"{api_base}/v1/transcriptions", + "api_base": join_api_base(api_base, "v1/transcriptions"), "atranscription": True, "complete_input_dict": SonioxAudioTranscriptionHandler._redact_body_for_logging(body), }, @@ -451,7 +452,7 @@ class SonioxAudioTranscriptionHandler: self._safe_log_pre_call(logging_obj, api_key, base_url, body) create_resp: Final = http_client.post( - url=f"{base_url}/v1/transcriptions", + url=join_api_base(base_url, "v1/transcriptions"), headers=auth_headers, json=body, timeout=timeout, @@ -472,7 +473,7 @@ class SonioxAudioTranscriptionHandler: ) transcript_resp: Final = http_client.get( - url=f"{base_url}/v1/transcriptions/{transcription_id}/transcript", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}/transcript"), headers=auth_headers, timeout=timeout, ) @@ -528,7 +529,7 @@ class SonioxAudioTranscriptionHandler: # `Authorization` header is fine; httpx sets multipart Content-Type. upload_headers: Final = {"Authorization": auth_headers["Authorization"]} resp: Final = http_client.post( - url=f"{base_url}/v1/files", + url=join_api_base(base_url, "v1/files"), headers=upload_headers, files=files, timeout=timeout, @@ -550,7 +551,7 @@ class SonioxAudioTranscriptionHandler: ) -> _TranscriptionMeta: for _ in range(max_attempts): resp = http_client.get( - url=f"{base_url}/v1/transcriptions/{transcription_id}", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}"), headers=auth_headers, timeout=timeout, ) @@ -593,7 +594,7 @@ class SonioxAudioTranscriptionHandler: if "transcription" in cleanup and transcription_id: try: http_client.delete( - url=f"{base_url}/v1/transcriptions/{transcription_id}", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}"), headers=auth_headers, timeout=timeout, ) @@ -605,7 +606,7 @@ class SonioxAudioTranscriptionHandler: if "file" in cleanup and file_id_to_cleanup: try: http_client.delete( - url=f"{base_url}/v1/files/{file_id_to_cleanup}", + url=join_api_base(base_url, f"v1/files/{file_id_to_cleanup}"), headers=auth_headers, timeout=timeout, ) @@ -685,7 +686,7 @@ class SonioxAudioTranscriptionHandler: self._safe_log_pre_call(logging_obj, api_key, base_url, body) create_resp: Final = await http_client.post( - url=f"{base_url}/v1/transcriptions", + url=join_api_base(base_url, "v1/transcriptions"), headers=auth_headers, json=body, timeout=timeout, @@ -706,7 +707,7 @@ class SonioxAudioTranscriptionHandler: ) transcript_resp: Final = await http_client.get( - url=f"{base_url}/v1/transcriptions/{transcription_id}/transcript", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}/transcript"), headers=auth_headers, timeout=timeout, ) @@ -761,7 +762,7 @@ class SonioxAudioTranscriptionHandler: } upload_headers: Final = {"Authorization": auth_headers["Authorization"]} resp: Final = await http_client.post( - url=f"{base_url}/v1/files", + url=join_api_base(base_url, "v1/files"), headers=upload_headers, files=files, timeout=timeout, @@ -783,7 +784,7 @@ class SonioxAudioTranscriptionHandler: ) -> _TranscriptionMeta: for _ in range(max_attempts): resp = await http_client.get( - url=f"{base_url}/v1/transcriptions/{transcription_id}", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}"), headers=auth_headers, timeout=timeout, ) @@ -826,7 +827,7 @@ class SonioxAudioTranscriptionHandler: if "transcription" in cleanup and transcription_id: try: await http_client.delete( - url=f"{base_url}/v1/transcriptions/{transcription_id}", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}"), headers=auth_headers, timeout=timeout, ) @@ -838,7 +839,7 @@ class SonioxAudioTranscriptionHandler: if "file" in cleanup and file_id_to_cleanup: try: await http_client.delete( - url=f"{base_url}/v1/files/{file_id_to_cleanup}", + url=join_api_base(base_url, f"v1/files/{file_id_to_cleanup}"), headers=auth_headers, timeout=timeout, ) diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 5d5fad6128e..974996c5e48 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -6,10 +6,12 @@ Handles transformation between OpenAI-compatible format and Stability AI API for API Reference: https://platform.stability.ai/docs/api-reference """ +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -30,6 +32,13 @@ else: LiteLLMLoggingObj = Any +class _StabilityImageEditResponse(TypedDict, total=False): + image: ReadOnly[str] + finish_reason: ReadOnly[str] + seed: ReadOnly[int] + errors: ReadOnly[Sequence[str]] + + class StabilityImageEditConfig(BaseImageEditConfig): """ Configuration for Stability AI image edit. @@ -259,7 +268,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): OpenAI expects: {"data": [{"b64_json": "base64..."}], "created": timestamp} """ try: - response_data: Final = raw_response.json() + response_data: Final[_StabilityImageEditResponse] = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error parsing Stability AI response: {e}", diff --git a/litellm/llms/stability/image_generation/transformation.py b/litellm/llms/stability/image_generation/transformation.py index 1856523ce49..a53fd087098 100644 --- a/litellm/llms/stability/image_generation/transformation.py +++ b/litellm/llms/stability/image_generation/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -131,11 +132,10 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the Stability AI API request. """ - base_url: str = api_base or get_secret_str("STABILITY_API_BASE") or self.DEFAULT_BASE_URL - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("STABILITY_API_BASE") or self.DEFAULT_BASE_URL endpoint: Final = self._get_model_endpoint(model) - return f"{base_url}{endpoint}" + return join_api_base(base_url, endpoint) def validate_environment( self, diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py index d0531e4c041..b46f974e0db 100644 --- a/litellm/llms/tavily/search/transformation.py +++ b/litellm/llms/tavily/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -90,11 +91,7 @@ class TavilySearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("TAVILY_API_BASE") or self.TAVILY_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("TAVILY_API_BASE") or self.TAVILY_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/tencent/chat/transformation.py b/litellm/llms/tencent/chat/transformation.py index 7e80b0012df..1f0bcd218d6 100644 --- a/litellm/llms/tencent/chat/transformation.py +++ b/litellm/llms/tencent/chat/transformation.py @@ -9,6 +9,7 @@ from typing import Final, TypedDict from typing_extensions import ReadOnly import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.secret_managers.main import get_secret_str from litellm.utils import supports_reasoning @@ -137,7 +138,4 @@ class TencentChatConfig(OpenAIGPTConfig): if api_base.endswith("/chat/completions"): return api_base - if not api_base.endswith("/v1"): - api_base = f"{api_base}/v1" - - return f"{api_base}/chat/completions" + return join_api_base(ensure_api_route(api_base, "v1"), "chat/completions") diff --git a/litellm/llms/tencent/messages/transformation.py b/litellm/llms/tencent/messages/transformation.py index c56ecaeb51c..30c3ecaf3a9 100644 --- a/litellm/llms/tencent/messages/transformation.py +++ b/litellm/llms/tencent/messages/transformation.py @@ -5,9 +5,8 @@ Tencent TokenHub exposes an Anthropic-compatible Messages API endpoint alongside its standard OpenAI-compatible chat completions endpoint. """ -from typing import Any - import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -47,7 +46,7 @@ class TencentAnthropicMessagesConfig(AnthropicMessagesConfig): self, headers: dict, model: str, - messages: list[Any], + messages: list[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, @@ -82,4 +81,4 @@ class TencentAnthropicMessagesConfig(AnthropicMessagesConfig): elif base_url.endswith("/v1"): base_url = base_url[: -len("/v1")] - return f"{base_url}/v1/messages" + return join_api_base(base_url, "v1/messages") diff --git a/litellm/llms/topaz/image_variations/transformation.py b/litellm/llms/topaz/image_variations/transformation.py index 8d03b73e94e..b3dded36f49 100644 --- a/litellm/llms/topaz/image_variations/transformation.py +++ b/litellm/llms/topaz/image_variations/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final from aiohttp import ClientResponse from httpx import Headers, Response +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import ( BaseLLMException, LiteLLMLoggingObj, @@ -39,8 +40,7 @@ class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - clean_api_base: Final = (api_base or "https://api.topazlabs.com").rstrip("/") - return f"{clean_api_base}/image/v1/enhance" + return join_api_base(api_base or "https://api.topazlabs.com", "image/v1/enhance") def map_openai_params( self, diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index b37bbd78f2b..fe0f96cb805 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -7,6 +7,7 @@ from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Final, Literal from httpx import Headers, Response +from typing_extensions import ReadOnly, TypedDict from litellm.constants import DEFAULT_MAX_TOKENS_FOR_TRITON from litellm.litellm_core_utils.prompt_templates.factory import prompt_factory @@ -32,6 +33,10 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +class _TritonGenerateResponse(TypedDict): + text_output: ReadOnly[str] + + class TritonConfig(BaseConfig): """ Base class for Triton configurations. @@ -220,7 +225,7 @@ class TritonGenerateConfig(TritonConfig): json_mode: bool | None = None, ) -> ModelResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[_TritonGenerateResponse] = raw_response.json() except Exception: raise TritonError(message=raw_response.text, status_code=raw_response.status_code) model_response.choices = [Choices(index=0, message=Message(content=raw_response_json["text_output"]))] diff --git a/litellm/llms/triton/embedding/transformation.py b/litellm/llms/triton/embedding/transformation.py index 8c6e0b832fc..3a735f741fd 100644 --- a/litellm/llms/triton/embedding/transformation.py +++ b/litellm/llms/triton/embedding/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Sequence from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.base_llm.chat.transformation import AllMessageValues, BaseLLMException from litellm.llms.base_llm.embedding.transformation import ( @@ -14,6 +16,16 @@ from litellm.utils import token_counter from ..common_utils import TritonError +class _TritonEmbeddingOutput(TypedDict): + shape: ReadOnly[Sequence[int]] + data: ReadOnly[Sequence[float]] + + +class _TritonEmbeddingResponse(TypedDict): + outputs: ReadOnly[Sequence[_TritonEmbeddingOutput]] + model_name: NotRequired[ReadOnly[str]] + + class TritonEmbeddingConfig(BaseEmbeddingConfig): """ Transformations for triton /embeddings endpoint (This is a trtllm model) @@ -79,7 +91,7 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict = {}, ) -> EmbeddingResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[_TritonEmbeddingResponse] = raw_response.json() except Exception: raise TritonError(message=raw_response.text, status_code=raw_response.status_code) @@ -134,7 +146,7 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig): return TritonError(message=error_message, status_code=status_code, headers=headers) @staticmethod - def split_embedding_by_shape(data: list[float], shape: list[int]) -> list[list[float]]: + def split_embedding_by_shape(data: Sequence[float], shape: Sequence[int]) -> list[Sequence[float]]: if len(shape) != 2: raise ValueError("Shape must be of length 2.") embedding_size: Final = shape[1] diff --git a/litellm/llms/vercel_ai_gateway/chat/transformation.py b/litellm/llms/vercel_ai_gateway/chat/transformation.py index 19c4c2659ca..cf633358a41 100644 --- a/litellm/llms/vercel_ai_gateway/chat/transformation.py +++ b/litellm/llms/vercel_ai_gateway/chat/transformation.py @@ -6,11 +6,14 @@ Calls done in OpenAI/openai.py as Vercel AI Gateway is openai-compatible. Docs: https://vercel.com/docs/ai-gateway """ +from collections.abc import Sequence from typing import Final import httpx +from typing_extensions import ReadOnly, TypedDict import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -19,6 +22,14 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig from ..common_utils import VercelAIGatewayException +class _VercelAIGatewayModel(TypedDict): + id: ReadOnly[str] + + +class _VercelAIGatewayModelList(TypedDict): + data: ReadOnly[Sequence[_VercelAIGatewayModel]] + + class VercelAIGatewayConfig(OpenAIGPTConfig): @property def custom_llm_provider(self) -> str | None: @@ -92,11 +103,11 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): if api_base is None: api_base = "https://ai-gateway.vercel.sh/v1" - models_url: Final = f"{api_base}/models" + models_url: Final = join_api_base(api_base, "models") response: Final = litellm.module_level_client.get(url=models_url) if response.status_code != 200: raise Exception(f"Failed to get models: {response.text}") - models: Final = response.json()["data"] - return [model["id"] for model in models] + models: Final[_VercelAIGatewayModelList] = response.json() + return [model["id"] for model in models["data"]] diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py index 9cfaaab89f0..01d7b7766c0 100644 --- a/litellm/llms/vercel_ai_gateway/embedding/transformation.py +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -77,12 +78,8 @@ class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): """ Get the complete URL for Vercel AI Gateway Embedding API endpoint. """ - if api_base: - api_base = api_base.rstrip("/") - else: - api_base = get_secret_str("VERCEL_AI_GATEWAY_API_BASE") or "https://ai-gateway.vercel.sh/v1" - - return f"{api_base}/embeddings" + base: Final = api_base or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") or "https://ai-gateway.vercel.sh/v1" + return join_api_base(base, "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py index f4db5eb110c..26de755d5ca 100644 --- a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py @@ -3,6 +3,7 @@ from collections.abc import Mapping, Sequence from typing import Final from httpx import Headers, Response +from pydantic import TypeAdapter import litellm from litellm.exceptions import UnsupportedParamsError @@ -10,6 +11,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( normalize_transcription_language_to_bcp47, process_audio_file, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -44,6 +46,7 @@ from litellm.types.utils import ( DEFAULT_GEMINI_TRANSCRIBE_LOCATION: Final = "global" AUDIO_MODALITY: Final = "AUDIO" +_DECODED_JSON: Final = TypeAdapter(object) class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase): @@ -129,10 +132,10 @@ class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexB project_id: Final = validate_vertex_transcription_project_id( self.safe_get_vertex_ai_project(vertex_params) or self._resolve_project_id_from_credentials(vertex_params) ) - base_url: Final = (api_base or get_vertex_base_url(location)).rstrip("/") + base_url: Final = api_base or get_vertex_base_url(location) bare_model: Final = model.removeprefix("vertex_ai/") model_path: Final = f"projects/{project_id}/locations/{location}/publishers/google/models/{bare_model}" - return f"{base_url}/v1/{model_path}:generateContent" + return join_api_base(base_url, f"v1/{model_path}:generateContent") def _resolve_project_id_from_credentials(self, litellm_params: Mapping[str, object]) -> str: vertex_params: Final = dict(litellm_params) @@ -176,7 +179,7 @@ class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexB raw_response: Response, ) -> TranscriptionResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except ValueError: raise VertexAIError( status_code=raw_response.status_code, diff --git a/litellm/llms/vertex_ai/audio_transcription/transformation.py b/litellm/llms/vertex_ai/audio_transcription/transformation.py index 5603be81095..4fdf93a6307 100644 --- a/litellm/llms/vertex_ai/audio_transcription/transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/transformation.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( normalize_transcription_language_to_bcp47, process_audio_file, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -126,8 +127,8 @@ class VertexAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase) project_id: Final = validate_vertex_transcription_project_id( self.safe_get_vertex_ai_project(litellm_params) or self._resolve_project_id_from_credentials(litellm_params) ) - base_url: Final = (api_base or f"https://{speech_to_text_host(location)}").rstrip("/") - return f"{base_url}/v2/projects/{project_id}/locations/{location}/recognizers/_:recognize" + base_url: Final = api_base or f"https://{speech_to_text_host(location)}" + return join_api_base(base_url, f"v2/projects/{project_id}/locations/{location}/recognizers/_:recognize") def _resolve_project_id_from_credentials(self, litellm_params: dict) -> str: _, project_id = self._ensure_access_token( diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index de5fe39cfc9..7c4238e3498 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -10,6 +10,7 @@ import litellm from litellm.litellm_core_utils.url_utils import ( async_safe_get, encode_url_path_segment, + join_api_base, safe_get, ) from litellm.llms.custom_httpx.http_handler import ( @@ -194,7 +195,7 @@ class VertexAIBatchPrediction(VertexLLM): api_base_path: Final = urlparse(api_base).path.rstrip("/") if api_base_path in ("/v1", "/v1beta1"): return graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url) - return api_base.rstrip("/") + urlparse(default_endpoint_url).path + return join_api_base(api_base, urlparse(default_endpoint_url).path) def _resolve_fine_tuned_endpoint_model( self, diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a8d02045740..99ef750b5fc 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -202,7 +202,7 @@ class VertexAIBatchTransformation: @classmethod def transform_vertex_ai_batch_list_response_to_openai_list_response( cls, response: dict[str, Any] - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """ Transforms Vertex AI batch list response into OpenAI-compatible list response. """ diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index f2da04a7db7..54c0fa97550 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -3,7 +3,7 @@ import json import os import time from collections.abc import Coroutine, Mapping -from typing import Any, Final +from typing import Final from urllib.parse import unquote import httpx @@ -197,7 +197,7 @@ class VertexAIFilesHandler(GCSBucketBase): timeout: float | httpx.Timeout, max_retries: int | None, litellm_params: dict | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: """ Download file content from GCS bucket for VertexAI files. Supports both sync and async operations. diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 176569f5d8d..410f29ef88a 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,9 +5,10 @@ import json import os import re import time -from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping +from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping, Sequence from contextlib import aclosing from dataclasses import dataclass +from functools import partial from types import MappingProxyType from typing import IO, Any, Final, TypedDict from urllib.parse import quote, unquote @@ -15,6 +16,7 @@ from urllib.parse import quote, unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from pydantic import TypeAdapter from typing_extensions import ReadOnly, Required import litellm @@ -30,11 +32,13 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, extract_file_metadata, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body @@ -93,6 +97,7 @@ _JSONL_NEWLINE: Final = b"\n" _BATCH_OUTPUT_FIRST_ROW_PEEK_LIMIT_BYTES: Final = 32 * 1024 * 1024 _PASSTHROUGH_MANAGED_GCS_PREFIX: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}passthrough/" _RAW_UPLOAD_CHUNK_BYTES: Final = 1024 * 1024 +_DECODED_JSON: Final = TypeAdapter(object) class _GcsObjectMetadataJson(TypedDict, total=False): @@ -824,7 +829,7 @@ def _iter_openai_jsonl_entries( def _load_openai_jsonl_row(lineno: int, line: str) -> object: try: - return json.loads(line) + return _DECODED_JSON.validate_python(json.loads(line)) except json.JSONDecodeError as e: raise VertexAIError( status_code=400, @@ -964,7 +969,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: Sequence[Mapping[str, object]], deployment_model: str | None = None, ) -> str: """ @@ -982,8 +987,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): `projects/../locations/../endpoints/` batch job model instead of a nonexistent publisher model. """ - first_entry: Final = openai_jsonl_content[0] - body_model: Final = first_entry["body"].get("model") if isinstance(first_entry.get("body"), dict) else None + body: Final = openai_jsonl_content[0].get("body") + body_fields: Final = as_str_mapping(body) + body_model: Final = body_fields.get("model") if body_fields is not None and isinstance(body, dict) else None jsonl_model: Final = body_model if isinstance(body_model, str) else "" raw_model: Final = deployment_model.removeprefix("vertex_ai/") if deployment_model else jsonl_model return _managed_batch_object_name(raw_model, passthrough=False) @@ -1081,12 +1087,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name: Final = encode_gcs_object_name_for_url(object_name) - endpoint: Final = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" + upload_path: Final = f"upload/storage/v1/b/{bucket_name}/o" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") - return f"{api_base}/{endpoint}" + return f"{join_api_base(api_base, upload_path)}?uploadType=media&name={encoded_object_name}" def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: return [] @@ -1103,6 +1109,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, openai_request_body: dict[str, Any], + deployment_drop_params: object = None, ) -> dict[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params @@ -1121,7 +1128,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): model=model, non_default_params=openai_request_body, optional_params={}, - drop_params=False, + drop_params=drops_unsupported_params(openai_request_body.get("drop_params", deployment_drop_params)), allow_integer_enums=False, ) return vertex_params @@ -1157,11 +1164,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): create_file_data=create_file_data, content_type=content_type, ): + deployment_drop_params: Final[object] = litellm_params.get("drop_params") return { "streaming_media_upload": StreamingMediaUploadConfig( body_stream=_OpenAIToVertexBatchUploadStream( file_data, - self._map_openai_to_vertex_params, + partial(self._map_openai_to_vertex_params, deployment_drop_params=deployment_drop_params), ), content_type="application/json", ) diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py index 40564d536a7..7cb52281a6a 100644 --- a/litellm/llms/vertex_ai/fine_tuning/handler.py +++ b/litellm/llms/vertex_ai/fine_tuning/handler.py @@ -27,6 +27,11 @@ from litellm.types.utils import LiteLLMFineTuningJob _JSON_VALUE: Final = TypeAdapter(object) +def _response_tuning_job(response: httpx.Response) -> ResponseTuningJob: + tuning_job: Final[ResponseTuningJob] = response.json() + return ResponseTuningJob(**tuning_job) + + class VertexFineTuningAPI(VertexLLM): """ Vertex methods to support for batches @@ -187,9 +192,7 @@ class VertexFineTuningAPI(VertexLLM): "got response from creating fine tuning job: %s", _JSON_VALUE.validate_python(response.json()) ) - vertex_response: Final = ResponseTuningJob( - **response.json(), - ) + vertex_response: Final = _response_tuning_job(response) verbose_logger.debug("vertex_response %s", vertex_response) open_ai_response: Final = self.convert_vertex_response_to_open_ai_response(vertex_response) @@ -272,9 +275,7 @@ class VertexFineTuningAPI(VertexLLM): verbose_logger.debug( "got response from creating fine tuning job: %s", _JSON_VALUE.validate_python(response.json()) ) - vertex_response: Final = ResponseTuningJob( - **response.json(), - ) + vertex_response: Final = _response_tuning_job(response) verbose_logger.debug("vertex_response %s", vertex_response) open_ai_response: Final = self.convert_vertex_response_to_open_ai_response(vertex_response) diff --git a/litellm/llms/vertex_ai/gemini/long_decode.py b/litellm/llms/vertex_ai/gemini/long_decode.py index 81d6c04d0e8..25aae70a853 100644 --- a/litellm/llms/vertex_ai/gemini/long_decode.py +++ b/litellm/llms/vertex_ai/gemini/long_decode.py @@ -2,17 +2,20 @@ import json from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from itertools import accumulate, chain -from typing import Final, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys +from typing import Final, Literal, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys import httpx +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.exceptions import Timeout +from litellm.litellm_core_utils.env_utils import get_env_int_or_none from litellm.types.llms.vertex_ai import ( Candidates, Citation, CitationMetadata, GenerateContentResponseBody, + GenerationConfig, GroundingMetadata, GroundingSegment, GroundingSupport, @@ -48,8 +51,14 @@ HOP_FAILURES: Final = (VertexAIError, httpx.HTTPError, Timeout) HOP_END_FINISH_REASON: Final = "CONTINUATION" CUT_OFF_FINISH_REASONS: Final = frozenset({"MAX_TOKENS", HOP_END_FINISH_REASON}) +DEFAULT_OUTPUT_BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" +_OUTPUT_LIMIT_KEYS: Final = frozenset({"max_output_tokens", "maxOutputTokens"}) +_REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET: Final = frozenset({"generation_config", "continuationToken"}) +_TEXT_ONLY_MODALITIES: Final = TypeAdapter(list[Literal["TEXT"]]) + _VERTEX_V1_PATH: Final = "/v1/projects/" _VERTEX_PATH_ACCEPTING_CONTINUATION_TOKEN: Final = "/v1beta1/projects/" +_CANDIDATE_COUNT_KEYS: Final = ("candidate_count", "candidateCount") LongDecodeHops: TypeAlias = tuple[GenerateContentResponseBody, ...] AsyncHopPost: TypeAlias = Callable[[RequestBody], Awaitable[httpx.Response]] @@ -76,15 +85,53 @@ def _int_or_none(value: object) -> int | None: return value if isinstance(value, int) else None +def _asks_for_one_candidate(generation_config: Mapping[str, object]) -> bool: + return all( + _int_or_none(generation_config[key] or 1) == 1 for key in _CANDIDATE_COUNT_KEYS if key in generation_config + ) + + def continuation_cap_of(request_body: RequestBody) -> int | None: generation_config: Final = request_body.get("generationConfig") if generation_config is None or "max_output_tokens" not in generation_config: return None - if "candidate_count" in generation_config and _int_or_none(generation_config["candidate_count"]) != 1: + if not _asks_for_one_candidate(generation_config): return None return _int_or_none(generation_config["max_output_tokens"]) +def _asks_for_text_only(generation_config: Mapping[str, object]) -> bool: + if "response_modalities" in generation_config: + return False + try: + _ = _TEXT_ONLY_MODALITIES.validate_python(generation_config.get("responseModalities", [])) + except ValidationError: + return False + return True + + +def with_default_output_budget( + request_body: RequestBody, model_max_output_tokens: Callable[[], int | None] +) -> RequestBody: + generation_config: Final = request_body.get("generationConfig", GenerationConfig()) + request_json: Final[Mapping[str, object]] = request_body + budget: Final = get_env_int_or_none(DEFAULT_OUTPUT_BUDGET_ENV) + if budget is None or budget <= 0: + return request_body + if not _REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET.isdisjoint(request_body): + return request_body + if not isinstance(request_json.get("generationConfig", {}), dict): + return request_body + if not _OUTPUT_LIMIT_KEYS.isdisjoint(generation_config): + return request_body + if not _asks_for_text_only(generation_config): + return request_body + model_max: Final = model_max_output_tokens() + if model_max is None: + return request_body + return {**request_body, "generationConfig": {**generation_config, "max_output_tokens": min(budget, model_max)}} + + def client_continuation_token(candidate: Candidates) -> str | None: if candidate.get("finishReason") not in CUT_OFF_FINISH_REASONS: return None @@ -285,7 +332,8 @@ def with_initial_prompt_tokens(usage: Usage, initial_prompt_tokens: int | None) def parse_body(response: httpx.Response) -> GenerateContentResponseBody: try: - return GenerateContentResponseBody(**response.json()) + body: Final[GenerateContentResponseBody] = response.json() + return GenerateContentResponseBody(**body) except Exception as e: raise VertexAIError( message=( diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index e5ce395031a..ed86934db63 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,16 +7,19 @@ Why separate file? Make it easy to see how transformation works import os import re from collections.abc import Callable, Mapping +from functools import partial from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote import httpx -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, PositiveInt, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_image_mime_type_from_url, ) @@ -65,6 +68,7 @@ from ..common_utils import ( get_supports_response_schema, get_supports_system_message, ) +from .long_decode import with_default_output_budget # VertexBase is imported only for type checking to avoid introducing a # module-load-time cyclic import to vertex_llm_base. The instance is lazily @@ -722,6 +726,9 @@ def gemini_convert_messages_with_history( from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig forward_function_call_id: Final = VertexGeminiConfig.forward_gemini_function_call_id(model or "") + supports_multimodal_function_responses: Final = model is None or ( + VertexGeminiConfig.supports_multimodal_function_responses(model) + ) try: while msg_i < len(messages): @@ -1025,6 +1032,7 @@ def gemini_convert_messages_with_history( messages[msg_i], last_message_with_tool_calls, forward_function_call_id=forward_function_call_id, + supports_multimodal_function_responses=supports_multimodal_function_responses, ) msg_i += 1 # Handle both single part and list of parts (for Computer Use with images) @@ -1304,13 +1312,75 @@ def sync_transform_request_body( vertex_auth_header=vertex_auth_header, ) - return transform_request_body( - messages=messages, - model=model, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - cached_content=cached_content, - optional_params=optional_params, + return _with_default_output_budget( + transform_request_body( + messages=messages, + model=model, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + cached_content=cached_content, + optional_params=optional_params, + ), + model, + custom_llm_provider, + litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict + ) + + +class _CostMapOutputModalities(TypedDict): + supported_output_modalities: NotRequired[ReadOnly[tuple[str, ...]]] + + +_COST_MAP_OUTPUT_MODALITIES: Final = TypeAdapter(_CostMapOutputModalities) + + +def _answers_text(cost_map_key: str) -> bool: + try: + row: Final = _COST_MAP_OUTPUT_MODALITIES.validate_python(litellm.model_cost.get(cost_map_key)) # pyright: ignore[reportUnknownMemberType] # cost map rows are untyped JSON + except ValidationError: + return False + return "text" in row.get("supported_output_modalities", ("text",)) + + +class _OutputLimit(TypedDict): + max_output_tokens: NotRequired[ReadOnly[PositiveInt | None]] + + +_OUTPUT_LIMIT: Final = TypeAdapter(_OutputLimit) + + +def _output_limit(fields: object) -> int | None: + try: + return _OUTPUT_LIMIT.validate_python(fields).get("max_output_tokens") + except ValidationError: + return None + + +def _catalog_capped_output_limit(cost_map_key: str, override: object) -> int | None: + limits: Final = ( + _output_limit(GetModelCostMap.loaded_model_cost_map().get(cost_map_key)), + _output_limit(override), + ) + return min((limit for limit in limits if limit is not None), default=None) + + +def _text_chat_model_max_output_tokens( + model: str, custom_llm_provider: str, deployment_model_info: object +) -> int | None: + info: Final = litellm.utils.lookup_model_info_for_registration(model, custom_llm_provider) + if isinstance(info, litellm.utils.ModelInfoLookupMiss) or info.get("mode") != "chat": + return None + if not _answers_text(info["key"]): + return None + return _catalog_capped_output_limit(info["key"], deployment_model_info or info) + + +def _with_default_output_budget( + request_body: RequestBody, model: str, custom_llm_provider: str, deployment_model_info: object +) -> RequestBody: + return with_default_output_budget( + request_body, + partial(_text_chat_model_max_output_tokens, model, custom_llm_provider, deployment_model_info), ) @@ -1380,21 +1450,8 @@ async def async_transform_request_body( messages, should_inline=should_inline_remote_media(custom_llm_provider) ) - if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): - # _transform_request_body may issue a sync httpx.get (up to 5s timeout) - # via _get_gcs_object_content_type to fetch GCS object metadata. Run the - # whole sync transformation on a worker thread so it does not block the - # async event loop. - return await asyncify(transform_request_body)( - messages=inlined_messages, - model=model, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - cached_content=cached_content, - optional_params=optional_params, - ) - - return transform_request_body( + transform: Final = partial( + transform_request_body, messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, @@ -1402,6 +1459,21 @@ async def async_transform_request_body( cached_content=cached_content, optional_params=optional_params, ) + # _transform_request_body may issue a sync httpx.get (up to 5s timeout) + # via _get_gcs_object_content_type to fetch GCS object metadata. Run the + # whole sync transformation on a worker thread so it does not block the + # async event loop. + request_body: Final[RequestBody] = ( + await asyncify(transform)() + if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages) + else transform() + ) + return _with_default_output_budget( + request_body, + model, + custom_llm_provider, + litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict + ) def _default_user_message_when_system_message_passed() -> ChatCompletionUserMessage: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 241b2d2a3a0..9acbc9c459c 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -12,7 +12,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args import httpx -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter import litellm from litellm import verbose_logger @@ -28,7 +28,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO, ) from litellm.exceptions import UnsupportedParamsError -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -42,6 +42,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, ) from litellm.types.llms.anthropic import AnthropicThinkingParam +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.gemini import BidiGenerateContentServerMessage from litellm.types.llms.openai import ( AllMessageValues, @@ -232,6 +233,274 @@ def _served_model_name(model_version: object) -> str | None: return model_version.split("@", 1)[0] +_GEMINI_SPEECH_CONFIG_KEY_MAP: Final[Mapping[str, str]] = MappingProxyType( + { + "multi_speaker_voice_config": "multiSpeakerVoiceConfig", + "speaker_voice_configs": "speakerVoiceConfigs", + "voice_config": "voiceConfig", + "prebuilt_voice_config": "prebuiltVoiceConfig", + "voice_name": "voiceName", + "replicated_voice_config": "replicatedVoiceConfig", + "mime_type": "mimeType", + "voice_sample_audio": "voiceSampleAudio", + "language_code": "languageCode", + "speaker_alias": "speakerAlias", + "speaker_id": "speakerId", + } +) + + +_SPEECH_CONFIG_ITEMS: Final = TypeAdapter(list[object]) + + +def _normalize_gemini_speech_config_item(value: Mapping[str, object]) -> Mapping[str, object]: + normalized: Final[dict[str, object]] = {} # mutable-ok: iterative traversal fills nested output without recursion + + pending: Final[ # mutable-ok: the explicit worklist must grow while avoiding recursive traversal + list[tuple[dict[str, object] | list[object], Mapping[str, object] | list[object]]] + ] = [(normalized, value)] + while pending: + target, source = pending.pop() + items = source.items() if isinstance(source, Mapping) else enumerate(source) + for key, item in items: + if isinstance(source, Mapping): + if not isinstance(key, str): + continue + normalized_key: str | int = _GEMINI_SPEECH_CONFIG_KEY_MAP.get(key, key) + else: + normalized_key = len(target) + + child: object + if (mapping_item := as_str_mapping(item)) is not None: + pending.append((child_mapping := dict[str, object](), mapping_item)) + child = child_mapping + elif isinstance(item, list): + pending.append((child_items := list[object](), _SPEECH_CONFIG_ITEMS.validate_python(item))) + child = child_items + else: + child = ( + OPENAI_TO_GEMINI_TTS_VOICES.get(item.lower(), item) + if normalized_key == "voiceName" and isinstance(item, str) + else item + ) + + if isinstance(target, dict): + assert isinstance(normalized_key, str) + target[normalized_key] = child + else: + target.append(child) + + return normalized + + +def _gemini_speaker_voice_config(item: object) -> object: + speaker_config: Final = as_str_mapping(item) + if speaker_config is None: + return item + speaker: Final = speaker_config.get("speakerAlias") + voice_name: Final = speaker_config.get("speakerId") + if not isinstance(speaker, str) or not isinstance(voice_name, str): + return item + return { + **{key: value for key, value in speaker_config.items() if key not in ("speakerAlias", "speakerId")}, + "speaker": speaker, + "voiceConfig": { + "prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name)} + }, + } + + +def _with_gemini_speaker_voice_configs(config: Mapping[str, object]) -> Mapping[str, object]: + multi_speaker_config: Final = as_str_mapping(config.get("multiSpeakerVoiceConfig")) + if multi_speaker_config is None: + return config + speaker_voice_configs: Final = multi_speaker_config.get("speakerVoiceConfigs") + if not isinstance(speaker_voice_configs, list): + return config + return { + **config, + "multiSpeakerVoiceConfig": { + **multi_speaker_config, + "speakerVoiceConfigs": [ + _gemini_speaker_voice_config(item) + for item in _SPEECH_CONFIG_ITEMS.validate_python(speaker_voice_configs) + ], + }, + } + + +_GEMINI_SPEECH_CONFIG_FIELDS: Final = frozenset({"voiceConfig", "multiSpeakerVoiceConfig", "languageCode"}) + +_GEMINI_VOICE_CONFIG_FIELDS: Final = MappingProxyType( + { + "prebuiltVoiceConfig": frozenset({"voiceName"}), + "replicatedVoiceConfig": frozenset({"mimeType", "voiceSampleAudio"}), + "voice": None, + } +) + +_GEMINI_SPEAKER_VOICE_CONFIG_FIELDS: Final = frozenset({"speaker", "voiceConfig"}) + + +def _known_leaf_fields(value: object, fields: frozenset[str] | None) -> object: + mapping: Final = as_str_mapping(value) + if fields is None or mapping is None: + return value + return {key: item for key, item in mapping.items() if key in fields} + + +def _known_voice_config_fields(value: object) -> object: + mapping: Final = as_str_mapping(value) + if mapping is None: + return value + return { + key: _known_leaf_fields(item, _GEMINI_VOICE_CONFIG_FIELDS[key]) + for key, item in mapping.items() + if key in _GEMINI_VOICE_CONFIG_FIELDS + } + + +def _known_speaker_voice_config_fields(value: object) -> object: + mapping: Final = as_str_mapping(value) + if mapping is None: + return value + return { + key: _known_voice_config_fields(item) if key == "voiceConfig" else item + for key, item in mapping.items() + if key in _GEMINI_SPEAKER_VOICE_CONFIG_FIELDS + } + + +def _known_speaker_voice_configs(value: object) -> object: + if not isinstance(value, list): + return value + return [_known_speaker_voice_config_fields(item) for item in _SPEECH_CONFIG_ITEMS.validate_python(value)] + + +def _known_multi_speaker_voice_config_fields(value: object) -> object: + mapping: Final = as_str_mapping(value) + if mapping is None: + return value + return {key: _known_speaker_voice_configs(item) for key, item in mapping.items() if key == "speakerVoiceConfigs"} + + +def _known_speech_config_field(key: str, value: object) -> object: + match key: + case "voiceConfig": + return _known_voice_config_fields(value) + case "multiSpeakerVoiceConfig": + return _known_multi_speaker_voice_config_fields(value) + case _: + return value + + +def _known_gemini_speech_fields( + config: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict + return { + key: _known_speech_config_field(key, item) + for key, item in config.items() + if key in _GEMINI_SPEECH_CONFIG_FIELDS + } + + +_GEMINI_TTS_VOICE_KEYS: Final = frozenset( + { + "format", + "name", + "voice", + "voiceName", + "voice_name", + "languageCode", + "language_code", + "modelName", + "model_name", + "speechConfig", + "speech_config", + "voiceConfig", + "voice_config", + "multiSpeakerVoiceConfig", + "multi_speaker_voice_config", + } +) + + +_NAMED_UNKNOWN_VOICE_KEYS: Final = 10 + +_NAMED_VOICE_KEY_CHARS: Final = 100 + + +def _gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object] | None: + camel_config: Final = as_str_mapping(value.get("speechConfig")) + return camel_config if camel_config is not None else as_str_mapping(value.get("speech_config")) + + +def named_voice_keys(keys: Sequence[str]) -> str: + named: Final = ", ".join(key[:_NAMED_VOICE_KEY_CHARS] for key in keys[:_NAMED_UNKNOWN_VOICE_KEYS]) + unnamed_count: Final = len(keys) - _NAMED_UNKNOWN_VOICE_KEYS + return f"{named}, and {unnamed_count} more" if unnamed_count > 0 else named + + +def reject_unknown_gemini_tts_voice_keys( + model: str, voice: Mapping[str, object], drop_params: bool, llm_provider: str | None +) -> None: + nested_config: Final = _gemini_speech_config_wrapper(voice) or {} + ignored_wrapper: Final = ( + ("speech_config",) + if as_str_mapping(voice.get("speechConfig")) is not None and voice.get("speech_config") is not None + else () + ) + unknown_keys: Final = sorted( + {*(key for key in (*voice, *nested_config) if key not in _GEMINI_TTS_VOICE_KEYS), *ignored_wrapper} + ) + if not unknown_keys or drops_unsupported_params(drop_params): + return + raise UnsupportedParamsError( + message=( + f"{model} does not support the voice keys {named_voice_keys(unknown_keys)}. Remove them, or set " + "drop_params to ignore them." + ), + model=model, + llm_provider=llm_provider, + ) + + +def normalize_gemini_speech_config( + value: Mapping[str, object], drop_unknown_fields: bool = False +) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict + speech_config: Final = {**_with_gemini_speaker_voice_configs(mapped_gemini_speech_config(value))} + return _known_gemini_speech_fields(speech_config) if drop_unknown_fields else speech_config + + +def _merge_gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object]: + nested_config: Final = _gemini_speech_config_wrapper(value) + if nested_config is None: + return value + return { + **{key: item for key, item in value.items() if key not in ("speechConfig", "speech_config")}, + **nested_config, + } + + +def mapped_gemini_speech_config( + value: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict + normalized: Final = _normalize_gemini_speech_config_item(_merge_gemini_speech_config_wrapper(value)) + speech_config: Final = {key: item for key, item in normalized.items() if key in _GEMINI_SPEECH_CONFIG_FIELDS} + voice_name: Final = next( + (normalized[key] for key in ("name", "voiceName", "voice") if isinstance(normalized.get(key), str)), + None, + ) + if "voiceConfig" in speech_config or "multiSpeakerVoiceConfig" in speech_config or not isinstance(voice_name, str): + return speech_config + return { + **speech_config, + "voiceConfig": { + "prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name)} + }, + } + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -420,6 +689,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) -> bool: return cls._forward_gemini_function_call_id(model) + @classmethod + def supports_multimodal_function_responses(cls, model: str) -> bool: + return cls.is_gemini_3_or_newer(model) + def _supports_penalty_parameters(self, model: str) -> bool: # Gemini 3 models do not support penalty parameters if VertexGeminiConfig._is_gemini_3_or_newer(model): @@ -1237,7 +1510,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): status_code=400, ) - def _map_audio_params(self, value: dict) -> dict: + def _map_audio_params( + self, value: Mapping[str, object], drop_unknown_fields: bool = False + ) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary """ Expected input: { @@ -1256,12 +1531,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): languageCode: "en-US", } """ - from litellm.types.llms.vertex_ai import ( - PrebuiltVoiceConfig, - SpeechConfig, - VoiceConfig, - ) - # Validate audio format - Gemini TTS only supports pcm16 audio_format: Final = value.get("format") if audio_format is not None and audio_format != "pcm16": @@ -1271,22 +1540,46 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): f"Please set audio format to 'pcm16'." ) - # Map OpenAI audio parameter to Gemini speech config - speech_config: Final[SpeechConfig] = {} + camel_speech_config: Final = value.get("speechConfig") + snake_speech_config: Final = value.get("speech_config") + speech_config_value: Final = ( + camel_speech_config if isinstance(camel_speech_config, Mapping) else snake_speech_config + ) + language_code: Final = value.get("language_code", value.get("languageCode")) + if isinstance(speech_config_value, Mapping): + return normalize_gemini_speech_config(value, drop_unknown_fields) - if "voice" in value: - voice: Final = value["voice"] - voice_name: Final = ( - OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice) if isinstance(voice, str) else voice - ) - prebuilt_voice_config: Final[PrebuiltVoiceConfig] = {"voiceName": voice_name} - voice_config: Final[VoiceConfig] = {"prebuiltVoiceConfig": prebuilt_voice_config} - speech_config["voiceConfig"] = voice_config + voice: Final = value.get("voice") + speech_config_without_language: Final[ # mutable-ok: provider request serialization requires a concrete dict + dict[str, object] + ] = ( + { + "voiceConfig": { + "prebuiltVoiceConfig": { + "voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice) + if isinstance(voice, str) + else voice, + } + } + } + if "voice" in value + else {} + ) - if "language_code" in value: - speech_config["languageCode"] = value["language_code"] + if language_code is not None: + return { + **speech_config_without_language, + "languageCode": language_code, + } - return cast(dict, speech_config) + return speech_config_without_language + + def _map_chat_audio_params( + self, model: str, value: Mapping[str, object], drop_params: bool + ) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary + if _gemini_speech_config_wrapper(value) is not None: + reject_unknown_gemini_tts_voice_keys(model, value, drop_params, self.custom_llm_provider) + return self._map_audio_params(value, drop_unknown_fields=drops_unsupported_params(drop_params)) def map_audio_params( self, @@ -1351,7 +1644,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif param == "n": optional_params["candidate_count"] = value elif param == "audio" and isinstance(value, dict): - optional_params["speechConfig"] = self._map_audio_params(value) + optional_params["speechConfig"] = self._map_chat_audio_params(model, value, drop_params) elif param == "stop": if isinstance(value, str): optional_params["stop_sequences"] = [value] @@ -2076,26 +2369,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif modality == "DOCUMENT": cached_text_tokens = (cached_text_tokens or 0) + token_count - ## Calculate non-cached tokens by subtracting cached from total (per modality) - ## This is necessary because promptTokensDetails includes both cached and non-cached tokens - ## See: https://github.com/BerriAI/litellm/issues/18750 - if cached_text_tokens is not None and prompt_text_tokens is not None: - # Explicit caching: subtract cached tokens per modality from cacheTokensDetails - prompt_text_tokens = prompt_text_tokens - cached_text_tokens - elif ( + if ( cached_tokens is not None and prompt_text_tokens is not None and cached_text_tokens is None and "cacheTokensDetails" not in usage_metadata ): - # Implicit caching: only cachedContentTokenCount is provided (no cacheTokensDetails) - # Subtract from text tokens since implicit caching is primarily for text content - # See: https://github.com/BerriAI/litellm/issues/16341 prompt_text_tokens = prompt_text_tokens - cached_tokens - if cached_audio_tokens is not None and prompt_audio_tokens is not None: - prompt_audio_tokens = prompt_audio_tokens - cached_audio_tokens - if cached_image_tokens is not None and prompt_image_tokens is not None: - prompt_image_tokens = prompt_image_tokens - cached_image_tokens if cached_video_tokens is not None and prompt_video_tokens is not None: prompt_video_tokens = prompt_video_tokens - cached_video_tokens @@ -2116,9 +2396,24 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): billable_text_tokens: Final = ( prompt_text_tokens if prompt_text_tokens is None else prompt_text_tokens + billable_tool_use_prompt_tokens ) + cached_tokens_details_fields: Final = MappingProxyType( + { + key: token_count + for key, token_count in ( + ("text_tokens", cached_text_tokens), + ("audio_tokens", cached_audio_tokens), + ("image_tokens", cached_image_tokens), + ) + if token_count is not None + } + ) + cached_tokens_details: Final = ( + CachedTokensDetails.model_validate(cached_tokens_details_fields) if cached_tokens_details_fields else None + ) prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=cached_tokens, + cached_tokens_details=cached_tokens_details, audio_tokens=prompt_audio_tokens, text_tokens=billable_text_tokens, image_tokens=prompt_image_tokens, diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 0a800497de1..f15c6393cb2 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -6,6 +6,7 @@ import json from typing import TYPE_CHECKING, Final, Literal import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.llms.custom_httpx.http_handler import ( @@ -35,6 +36,11 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +class _GeminiFileResource(TypedDict, total=False): + mimeType: ReadOnly[str] + uri: ReadOnly[str] + + class GoogleBatchEmbeddings(VertexLLM): @staticmethod def _flatten_and_detect_file_refs( @@ -74,7 +80,7 @@ class GoogleBatchEmbeddings(VertexLLM): if response.status_code != 200: raise Exception(f"Error fetching file {element}: {response.status_code} {response.text}") - file_data = response.json() + file_data: _GeminiFileResource = response.json() resolved_files[element] = { "mime_type": file_data.get("mimeType", ""), "uri": file_data.get("uri", element), @@ -111,7 +117,7 @@ class GoogleBatchEmbeddings(VertexLLM): if response.status_code != 200: raise Exception(f"Error fetching file {element}: {response.status_code} {response.text}") - file_data = response.json() + file_data: _GeminiFileResource = response.json() resolved_files[element] = { "mime_type": file_data.get("mimeType", ""), "uri": file_data.get("uri", element), diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 725a7f39917..6d071363823 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -190,7 +190,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): payload: Final = json.dumps(request_body) empty_files: Final = cast(RequestFiles, []) - return cast(tuple[dict[str, Any], RequestFiles | None], (payload, empty_files)) + return cast(tuple[dict[str, object], RequestFiles | None], (payload, empty_files)) def transform_image_edit_response( self, diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index f916d494aad..deaa32f7308 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -9,9 +9,11 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles from pydantic import ConfigDict, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM @@ -33,6 +35,10 @@ _JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object] ) +class _VertexImagenPredictResponse(TypedDict, total=False): + predictions: ReadOnly[Sequence[object]] + + class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """ Vertex AI Imagen Image Edit Configuration @@ -143,12 +149,12 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): if model.startswith("vertex_ai/"): model_name = model.replace("vertex_ai/", "") - if api_base: - base_url = api_base.rstrip("/") - else: - base_url = get_vertex_base_url(vertex_location) + base_url: Final = api_base or get_vertex_base_url(vertex_location) - return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:predict" + return join_api_base( + base_url, + f"v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:predict", + ) def transform_image_edit_request( self, @@ -206,7 +212,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): ) -> ImageResponse: model_response: Final = ImageResponse() try: - response_json: Final = raw_response.json() + response_json: Final[_VertexImagenPredictResponse] = raw_response.json() except Exception as exc: raise self.get_error_class( error_message=f"Error transforming image edit response: {exc}", diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index 7ba7b060991..ec915e10fcd 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -168,7 +168,7 @@ class VertexImageGeneration(VertexLLM): if response.status_code != 200: raise Exception(f"Error: {response.status_code} {response.text}") - json_response: Final = response.json() + json_response: Final[Mapping[str, object]] = response.json() return self.process_image_generation_response(json_response, model_response, model) async def aimage_generation( @@ -272,7 +272,7 @@ class VertexImageGeneration(VertexLLM): if response.status_code != 200: raise Exception(f"Error: {response.status_code} {response.text}") - json_response: Final = response.json() + json_response: Final[Mapping[str, object]] = response.json() return self.process_image_generation_response(json_response, model_response, model) def is_image_generation_response(self, json_response: dict[str, Any]) -> bool: diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 916e25a027e..f6df772dd56 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -33,6 +33,7 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( @@ -299,7 +300,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): Transform Gemini image generation response to litellm ImageResponse format """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py index 36965fe666f..b8e13059917 100644 --- a/litellm/llms/vertex_ai/interactions/transformation.py +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -2,7 +2,7 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Final -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig from litellm.llms.vertex_ai.common_utils import validate_vertex_location from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -22,9 +22,9 @@ class VertexInteractionsTarget: @property def collection_url(self) -> str: - return ( - f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}" - f"/projects/{self.project_id}/locations/{self.location}/interactions" + return join_api_base( + self.base_url, + f"{VERTEX_INTERACTIONS_API_VERSION}/projects/{self.project_id}/locations/{self.location}/interactions", ) def interaction_url(self, interaction_id: str) -> str: diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index da77fa338db..c282ae22ea6 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -208,7 +208,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): if raw_response.status_code != 200: raise Exception(f"Error: {raw_response.status_code} {raw_response.text}") - _json_response: Final = raw_response.json() + _json_response: Final[MultimodalPredictions] = raw_response.json() if "predictions" not in _json_response: raise InternalServerError( message=f"embedding response does not contain 'predictions', got {_json_response}", diff --git a/litellm/llms/vertex_ai/rag_engine/transformation.py b/litellm/llms/vertex_ai/rag_engine/transformation.py index a34ef519578..61cb083fa4c 100644 --- a/litellm/llms/vertex_ai/rag_engine/transformation.py +++ b/litellm/llms/vertex_ai/rag_engine/transformation.py @@ -4,7 +4,10 @@ Transformation utilities for Vertex AI RAG Engine. Handles transforming LiteLLM's unified formats to Vertex AI RAG Engine API format. """ -from typing import Any, Final +from collections.abc import Mapping +from typing import Final + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE @@ -13,6 +16,15 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.rag import RAGChunkingStrategy +class VertexRAGChunkingConfig(TypedDict): + chunk_size: ReadOnly[int] + chunk_overlap: ReadOnly[int] + + +class VertexRAGTransformationConfig(TypedDict): + chunking_config: ReadOnly[VertexRAGChunkingConfig] + + class VertexAIRAGTransformation(VertexBase): """ Transformation class for Vertex AI RAG Engine API. @@ -55,7 +67,7 @@ class VertexAIRAGTransformation(VertexBase): def transform_chunking_strategy_to_vertex_format( self, chunking_strategy: RAGChunkingStrategy | None, - ) -> dict[str, Any]: + ) -> VertexRAGTransformationConfig: """ Transform LiteLLM's unified chunking_strategy to Vertex AI RAG format. @@ -105,7 +117,7 @@ class VertexAIRAGTransformation(VertexBase): self, gcs_uri: str, chunking_strategy: RAGChunkingStrategy | None = None, - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """ Build the request payload for importing RAG files. diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index 55f1cf58cbe..13f547b3b92 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -21,6 +21,7 @@ from collections.abc import Awaitable, Callable from typing import Final from litellm import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( VertexChirpRealtimeConfig, @@ -68,9 +69,8 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): """ if api_base: # Allow callers to supply a fully-qualified wss:// base URL. - base = api_base.rstrip("/") - base = base.replace("https://", "wss://").replace("http://", "ws://") - return f"{base}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" + base: Final = api_base.replace("https://", "wss://").replace("http://", "ws://") + return join_api_base(base, "ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent") host: Final = get_vertex_base_url(self._location).removeprefix("https://") return f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index dce4d2f2a87..398471a7c9b 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -6,10 +6,11 @@ Why separate file? Make it easy to see how transformation works import math import uuid -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -27,6 +28,17 @@ from litellm.types.rerank import ( ) +class _VertexRankedRecord(TypedDict): + id: ReadOnly[str] + score: NotRequired[ReadOnly[float]] + title: NotRequired[ReadOnly[str]] + content: NotRequired[ReadOnly[str]] + + +class _VertexRankResponse(TypedDict, total=False): + records: ReadOnly[Sequence[_VertexRankedRecord]] + + class VertexAIRerankConfig(BaseRerankConfig, VertexBase): """ Configuration for Vertex AI Discovery Engine Rerank API @@ -172,7 +184,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): Transform Vertex AI Discovery Engine response to Cohere format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[_VertexRankResponse] = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse response: {e}") diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 2fae2a9de1c..c0082bc5ca6 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -6,20 +6,25 @@ Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/s """ import base64 -from collections.abc import Coroutine, Mapping +import math +from collections.abc import Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypeAlias, Union +from typing import TYPE_CHECKING, Any, ClassVar, Final, TypeAlias, Union import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm +from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.audio_utils.utils import ( DEFAULT_SPEECH_MEDIA_TYPE, + calculate_request_duration, speech_media_type_from_audio_bytes, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -33,6 +38,7 @@ from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.llms.vertex_ai_text_to_speech import ( VertexTextToSpeechAudioConfig, VertexTextToSpeechInput, + VertexTextToSpeechSpeakerVoiceConfig, VertexTextToSpeechVoice, ) @@ -49,6 +55,166 @@ _LyriaVoice: TypeAlias = str | dict | None _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_ARRAY: Final = TypeAdapter(tuple[object, ...], config=ConfigDict(hide_input_in_errors=True)) _STR: Final = TypeAdapter(str, config=ConfigDict(hide_input_in_errors=True)) +_GEMINI_CLOUD_TTS_LOCATIONS: Final = MappingProxyType( + { + **dict.fromkeys( + ("us", "us-central1", "us-east1", "us-east4", "us-east5", "us-south1", "us-west1", "us-west4"), "us" + ), + **dict.fromkeys( + ("eu", "europe-central2", "europe-north1", "europe-southwest1", "europe-west1", "europe-west4"), "eu" + ), + "northamerica-northeast1": "northamerica-northeast1", + } +) + + +_DEFAULT_GEMINI_TTS_VOICE: Final = "Kore" + + +class _GeminiCloudTtsInput(BaseModel): + model_config = ConfigDict(frozen=True) + + text: str | None = None + ssml: str | None = None + prompt: str | None = None + + +class _GeminiCloudTtsAudioConfig(BaseModel): + model_config = ConfigDict(frozen=True) + + audioEncoding: str = "" + sampleRateHertz: int | None = None + + +class _GeminiCloudTtsRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + input: _GeminiCloudTtsInput + audioConfig: _GeminiCloudTtsAudioConfig + + +def _fallback_gemini_tts_audio_duration(audio: bytes, encoding: str, sample_rate: int) -> float | None: + if encoding == "PCM": + return len(audio) / (2 * sample_rate) if sample_rate > 0 else None + if encoding in ("ALAW", "MULAW") and not (audio[:4] == b"RIFF" and audio[8:12] == b"WAVE"): + return len(audio) / sample_rate if sample_rate > 0 else None + + if encoding in ("LINEAR16", "ALAW", "MULAW") and audio[:4] == b"RIFF" and audio[8:12] == b"WAVE": + format_offset: Final = audio.find(b"fmt ", 12) + data_offset: Final = audio.find(b"data", 12) + if format_offset >= 0 and data_offset >= 0 and format_offset + 20 <= len(audio): + byte_rate: Final = int.from_bytes(audio[format_offset + 16 : format_offset + 20], "little") + data_size: Final = int.from_bytes(audio[data_offset + 4 : data_offset + 8], "little") + if byte_rate > 0 and data_size <= len(audio) - data_offset - 8: + return data_size / byte_rate + + if encoding == "MP3" and speech_media_type_from_audio_bytes(audio) == "audio/mpeg": + return len(audio) / 4000 + + if encoding == "OGG_OPUS" and audio[:4] == b"OggS": + opus_header: Final = audio.find(b"OpusHead") + last_page: Final = audio.rfind(b"OggS") + if opus_header >= 0 and opus_header + 12 <= len(audio) and last_page + 14 <= len(audio): + pre_skip: Final = int.from_bytes(audio[opus_header + 10 : opus_header + 12], "little") + granule: Final = int.from_bytes(audio[last_page + 6 : last_page + 14], "little") + if pre_skip <= granule < 2**64 - 1: + return (granule - pre_skip) / 48000 + + return None + + +_SPEAKER_CONFIG_ITEMS: Final = TypeAdapter(list[object]) + +_CloudTtsVoiceFields: TypeAlias = "Mapping[str, _CloudTtsVoiceFields]" + +_CLOUD_TTS_VOICE_LEAF: Final[_CloudTtsVoiceFields] = MappingProxyType({}) + +_CLOUD_TTS_PREBUILT_VOICE_FIELDS: Final[_CloudTtsVoiceFields] = MappingProxyType( + {"prebuiltVoiceConfig": MappingProxyType({"voiceName": _CLOUD_TTS_VOICE_LEAF})} +) + +_CLOUD_TTS_SPEECH_CONFIG_FIELDS: Final[_CloudTtsVoiceFields] = MappingProxyType( + { + "languageCode": _CLOUD_TTS_VOICE_LEAF, + "voiceConfig": _CLOUD_TTS_PREBUILT_VOICE_FIELDS, + "multiSpeakerVoiceConfig": MappingProxyType( + { + "speakerVoiceConfigs": MappingProxyType( + { + "speaker": _CLOUD_TTS_VOICE_LEAF, + "speakerAlias": _CLOUD_TTS_VOICE_LEAF, + "speakerId": _CLOUD_TTS_VOICE_LEAF, + "voiceConfig": _CLOUD_TTS_PREBUILT_VOICE_FIELDS, + } + ) + } + ), + } +) + + +_CLOUD_TTS_LIST_FIELDS: Final = frozenset({"speakerVoiceConfigs"}) + +_GEMINI_TTS_VOICE_NAME_KEYS: Final = ("name", "voiceName", "voice_name", "voice") + + +def _fields_cloud_tts_cannot_send( + config: Mapping[str, object], fields: _CloudTtsVoiceFields, path: str +) -> Iterator[str]: + for key, value in config.items(): + yield from _field_cloud_tts_cannot_send(f"{path}{key}", value, fields.get(key), key in _CLOUD_TTS_LIST_FIELDS) + + +def _field_cloud_tts_cannot_send( + path: str, value: object, fields: "_CloudTtsVoiceFields | None", is_list_field: bool +) -> Iterator[str]: + if fields is None: + yield path + return + if value is None: + return + if not fields: + yield from () if isinstance(value, str) else (path,) + return + items: Final = ( + _SPEAKER_CONFIG_ITEMS.validate_python(value) if is_list_field and isinstance(value, list) else (value,) + ) + for item in items: + mapping = as_str_mapping(item) + yield from (path,) if mapping is None else _fields_cloud_tts_cannot_send(mapping, fields, f"{path}.") + + +def _non_string_voice_names(config: Mapping[str, object], path: str) -> Iterator[str]: + return ( + f"{path}{key}" + for key in _GEMINI_TTS_VOICE_NAME_KEYS + if config.get(key) is not None and not isinstance(config.get(key), str) + ) + + +def _non_mapping_speech_config_wrappers(voice: Mapping[str, object]) -> Iterator[str]: + return ( + key + for key in ("speechConfig", "speech_config") + if voice.get(key) is not None and as_str_mapping(voice.get(key)) is None + ) + + +def _speaker_entries(config: Mapping[str, object]) -> Iterator[Mapping[str, object]]: + multi_speaker_config: Final = as_str_mapping(config.get("multiSpeakerVoiceConfig")) + speakers: Final = multi_speaker_config.get("speakerVoiceConfigs") if multi_speaker_config is not None else None + items: Final = _SPEAKER_CONFIG_ITEMS.validate_python(speakers) if isinstance(speakers, list) else () + return (entry for item in items if (entry := as_str_mapping(item)) is not None) + + +def _prebuilt_voices_replaced_by_speaker_id(config: Mapping[str, object]) -> Iterator[str]: + return ( + "multiSpeakerVoiceConfig.speakerVoiceConfigs.voiceConfig.prebuiltVoiceConfig" + for entry in _speaker_entries(config) + if entry.get("speakerId") is not None + and (voice_config := as_str_mapping(entry.get("voiceConfig"))) is not None + and voice_config.get("prebuiltVoiceConfig") is not None + ) class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): @@ -81,14 +247,19 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): } # Response format mappings from OpenAI to Google Cloud audio encoding - FORMAT_MAPPINGS = { - "mp3": "MP3", - "opus": "OGG_OPUS", - "aac": "MP3", # Google doesn't have AAC, use MP3 - "flac": "FLAC", - "wav": "LINEAR16", - "pcm": "LINEAR16", - } + FORMAT_MAPPINGS: ClassVar[Mapping[str, str]] = MappingProxyType( + { + "mp3": "MP3", + "opus": "OGG_OPUS", + "aac": "MP3", # Google doesn't have AAC, use MP3 + "flac": "FLAC", + "wav": "LINEAR16", + "pcm": "LINEAR16", + } + ) + GEMINI_FORMAT_MAPPINGS: ClassVar[Mapping[str, str]] = MappingProxyType( + {"mp3": "MP3", "opus": "OGG_OPUS", "alaw": "ALAW", "mulaw": "MULAW"} + ) def __init__(self) -> None: BaseTextToSpeechConfig.__init__(self) @@ -96,7 +267,9 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): def _map_voice_to_vertex_format( self, - voice: str | dict | None, + voice: str | dict[str, object] | None, + model: str | None = None, + drop_params: bool = False, ) -> tuple[str | None, dict | None]: """ Map voice to Vertex AI format. @@ -111,6 +284,9 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): - voice_str: Original string voice (for interface compatibility) - voice_dict: Vertex AI format dict with languageCode and name """ + if model is not None and self._is_gemini_tts_model(model): + return self._map_gemini_tts_voice_to_vertex_format(model=model, voice=voice, drop_params=drop_params) + if voice is None: return None, None @@ -142,11 +318,225 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): return voice_str, voice_dict + @staticmethod + def _is_gemini_tts_model(model: str) -> bool: + from litellm.utils import is_gemini_tts_model + + return is_gemini_tts_model(model) + + @staticmethod + def _get_str_value( + source: Mapping[str, object], + *keys: str, + ) -> str | None: + for key in keys: + value = source.get(key) + if isinstance(value, str): + return value + return None + + @staticmethod + def _get_dict_value( + source: Mapping[str, object], + *keys: str, + ) -> Mapping[str, object] | None: + for key in keys: + value = source.get(key) + if (mapping := as_str_mapping(value)) is not None: + return mapping + return None + + @staticmethod + def _resolve_gemini_tts_voice_name(model: str, voice_name: str, drop_params: bool) -> str: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import OPENAI_TO_GEMINI_TTS_VOICES + + gemini_voice_name: Final = OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name) + if "-" not in gemini_voice_name: + return gemini_voice_name + if drop_params: + return _DEFAULT_GEMINI_TTS_VOICE + raise UnsupportedParamsError( + message=( + f"{model} speaks Gemini voices such as {_DEFAULT_GEMINI_TTS_VOICE}, and {voice_name} is a Cloud " + "Text-to-Speech voice. Use a Gemini voice, or set drop_params to use the default Gemini voice." + ), + model=model, + llm_provider="vertex_ai", + ) + + @staticmethod + def _reject_malformed_gemini_tts_voice(model: str, drop_params: bool) -> None: + if drop_params: + return + raise UnsupportedParamsError( + message=( + f"{model} needs voice to be a Gemini voice name such as {_DEFAULT_GEMINI_TTS_VOICE} or a speech " + "config object. Send one, or set drop_params to use the default Gemini voice." + ), + model=model, + llm_provider="vertex_ai", + ) + + @staticmethod + def _reject_voice_fields_cloud_tts_cannot_send(model: str, voice: Mapping[str, object], drop_params: bool) -> None: + if drop_params: + return + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + mapped_gemini_speech_config, + named_voice_keys, + ) + + speech_config: Final = VertexAITextToSpeechConfig._get_dict_value(voice, "speechConfig", "speech_config") + mapped_config: Final = mapped_gemini_speech_config(voice) + unsendable: Final = sorted( + { + *_fields_cloud_tts_cannot_send(mapped_config, _CLOUD_TTS_SPEECH_CONFIG_FIELDS, ""), + *_prebuilt_voices_replaced_by_speaker_id(mapped_config), + *_non_mapping_speech_config_wrappers(voice), + *_non_string_voice_names(voice, ""), + *_non_string_voice_names(speech_config or {}, "speechConfig."), + } + ) + if not unsendable: + return + raise UnsupportedParamsError( + message=( + f"{model} cannot send the voice fields {named_voice_keys(unsendable)} as mp3, opus, alaw, or mulaw. " + "Request wav or pcm, remove them, or set drop_params to ignore them." + ), + model=model, + llm_provider="vertex_ai", + ) + + @staticmethod + def _reject_malformed_gemini_tts_speaker(model: str, drop_params: bool) -> None: + if drop_params: + return + raise UnsupportedParamsError( + message=( + f"{model} needs each entry of speaker_voice_configs to name a speaker and a voice, as speaker with " + "voice_config.prebuilt_voice_config.voice_name, or speakerAlias with speakerId." + ), + model=model, + llm_provider="vertex_ai", + ) + + def _extract_gemini_tts_speaker_configs( + self, + model: str, + speech_config: Mapping[str, object], + drop_params: bool, + ) -> Sequence[VertexTextToSpeechSpeakerVoiceConfig]: + multi_speaker_config: Final = self._get_dict_value(speech_config, "multiSpeakerVoiceConfig") + if multi_speaker_config is None: + return [] + raw_speaker_configs: Final = multi_speaker_config.get("speakerVoiceConfigs", []) + try: + validated_speaker_configs: Final = _SPEAKER_CONFIG_ITEMS.validate_python(raw_speaker_configs, strict=True) + except ValidationError: + self._reject_malformed_gemini_tts_speaker(model, drop_params) + return [] + speaker_configs: Final[ # mutable-ok: validated payloads are accumulated for serialization + list[VertexTextToSpeechSpeakerVoiceConfig] + ] = [] + for raw_item in validated_speaker_configs: + if (raw_config := as_str_mapping(raw_item)) is None: + self._reject_malformed_gemini_tts_speaker(model, drop_params) + continue + speaker_alias = self._get_str_value(raw_config, "speakerAlias", "speaker") + speaker_id = self._get_str_value(raw_config, "speakerId") or self._prebuilt_voice_name(raw_config) + if speaker_alias is None or speaker_id is None: + self._reject_malformed_gemini_tts_speaker(model, drop_params) + continue + speaker_configs.append( + { + "speakerAlias": speaker_alias, + "speakerId": self._resolve_gemini_tts_voice_name(model, speaker_id, drop_params), + } + ) + return speaker_configs + + def _prebuilt_voice_name(self, config: Mapping[str, object]) -> str | None: + voice_config: Final = self._get_dict_value(config, "voiceConfig") + prebuilt_voice_config: Final = ( + self._get_dict_value(voice_config, "prebuiltVoiceConfig") if voice_config is not None else None + ) + if prebuilt_voice_config is None: + return None + return self._get_str_value(prebuilt_voice_config, "voiceName") + + def _map_gemini_tts_voice_to_vertex_format( + self, + model: str, + voice: object, + drop_params: bool = False, + ) -> tuple[str | None, dict[str, object]]: # mutable-ok: provider request serialization requires a concrete dict + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + mapped_gemini_speech_config, + reject_unknown_gemini_tts_voice_keys, + ) + + default_voice: Final[dict[str, object]] = { + "languageCode": self.DEFAULT_LANGUAGE_CODE, + "modelName": model, + "name": _DEFAULT_GEMINI_TTS_VOICE, + } + if voice is None: + return None, default_voice + + if isinstance(voice, str) and voice: + return voice, { + "languageCode": self.DEFAULT_LANGUAGE_CODE, + "modelName": model, + "name": self._resolve_gemini_tts_voice_name(model, voice, drop_params), + } + + voice_config: Final = as_str_mapping(voice) + if voice_config is None: + self._reject_malformed_gemini_tts_voice(model, drop_params) + return None, default_voice + reject_unknown_gemini_tts_voice_keys(model, voice_config, drop_params, "vertex_ai") + + speech_config: Final = mapped_gemini_speech_config(voice_config) + language_code: Final = self._get_str_value(speech_config, "languageCode") or self.DEFAULT_LANGUAGE_CODE + model_name: Final = model + speaker_configs: Final = self._extract_gemini_tts_speaker_configs(model, speech_config, drop_params) + self._reject_voice_fields_cloud_tts_cannot_send(model, voice_config, drop_params) + if speaker_configs: + return None, { + "languageCode": language_code, + "modelName": model_name, + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": speaker_configs, + }, + } + voice_name: Final = self._prebuilt_voice_name(speech_config) + if voice_name is not None: + return None, { + "languageCode": language_code, + "modelName": model_name, + "name": self._resolve_gemini_tts_voice_name(model, voice_name, drop_params), + } + return None, { + "languageCode": language_code, + "modelName": model_name, + "name": _DEFAULT_GEMINI_TTS_VOICE, + } + + @staticmethod + def _dispatch_voice_name(voice: str | Mapping[str, object] | None) -> str | None: + if isinstance(voice, str): + return voice + if not isinstance(voice, Mapping): + return None + name: Final = voice.get("name") + return name if isinstance(name, str) else None + def dispatch_text_to_speech( self, model: str, input: str, - voice: str | dict | None, + voice: str | dict[str, object] | None, optional_params: dict[str, object], litellm_params_dict: dict[str, object], logging_obj: "LiteLLMLoggingObj", @@ -175,14 +565,17 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): vertex_project: Final = self.safe_get_vertex_ai_project(litellm_params_dict) vertex_location: Final = self.safe_get_vertex_ai_location(litellm_params_dict) - # Convert voice to string if it's a dict (extract name) - # Actual voice mapping happens in map_openai_params - voice_str: str | None = None - if isinstance(voice, str): - voice_str = voice - elif isinstance(voice, dict): - # Extract voice name from dict if needed - voice_str = voice.get("name") if voice else None + mapped_voice, mapped_params = ( + self.map_openai_params( + model=model, + voice=voice, + optional_params=optional_params, + drop_params=drops_unsupported_params(kwargs.get("drop_params")), + kwargs=kwargs, + ) + if self._is_gemini_tts_model(model) + else (self._dispatch_voice_name(voice), optional_params) + ) # Store credentials in litellm_params for use in transform methods litellm_params_dict.update( @@ -198,9 +591,11 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): response: Final = base_llm_http_handler.text_to_speech_handler( model=model, input=input, - voice=voice_str, - text_to_speech_provider_config=self, - text_to_speech_optional_params=optional_params, + voice=mapped_voice, + text_to_speech_provider_config=( + VertexAIGeminiCloudTextToSpeechConfig() if self._is_gemini_tts_model(model) else self + ), + text_to_speech_optional_params=mapped_params, custom_llm_provider="vertex_ai", litellm_params=litellm_params_dict, logging_obj=logging_obj, @@ -219,13 +614,13 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): Note: Vertex AI also supports additional parameters like audioConfig which can be passed but are not part of the OpenAI spec """ - return ["voice", "response_format", "speed"] + return ["voice", "response_format", "speed", *(["instructions"] if self._is_gemini_tts_model(model) else [])] def map_openai_params( self, model: str, optional_params: dict, - voice: str | dict | None = None, + voice: str | dict[str, object] | None = None, drop_params: bool = False, kwargs: dict = {}, ) -> tuple[str | None, dict]: @@ -244,19 +639,28 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): Tuple of (mapped_voice_str, mapped_params) """ mapped_params: Final[dict[str, object]] = {} + openai_params: Final = _JSON_OBJECT.validate_python(optional_params) + drops_params: Final = drop_params or litellm.drop_params is True ########################################################## # Map voice using helper ########################################################## - mapped_voice_str, voice_dict = self._map_voice_to_vertex_format(voice) + mapped_voice_str, voice_dict = self._map_voice_to_vertex_format( + voice=voice, + model=model, + drop_params=drops_params, + ) if voice_dict is not None: mapped_params["vertex_voice_dict"] = voice_dict # Map response format if "response_format" in optional_params: format_name: Final = optional_params["response_format"] - if format_name in self.FORMAT_MAPPINGS: - mapped_params["audioEncoding"] = self.FORMAT_MAPPINGS[format_name] + format_mappings: Final = ( + self.GEMINI_FORMAT_MAPPINGS if self._is_gemini_tts_model(model) else self.FORMAT_MAPPINGS + ) + if format_name in format_mappings: + mapped_params["audioEncoding"] = format_mappings[format_name] else: # Try to use it directly as Google Cloud format mapped_params["audioEncoding"] = format_name @@ -265,10 +669,19 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): mapped_params["audioEncoding"] = self.DEFAULT_AUDIO_ENCODING # Map speed (OpenAI: 0.25-4.0, Vertex AI: speakingRate 0.25-4.0) - if "speed" in optional_params: - speed: Final = optional_params["speed"] - if speed is not None: - mapped_params["speakingRate"] = str(speed) + speed: Final = openai_params.get("speed") + if speed is not None: + mapped_params["speakingRate"] = str(speed) + + instructions: Final = openai_params.get("instructions") + if isinstance(instructions, str) and self._is_gemini_tts_model(model): + mapped_params["prompt"] = instructions + elif instructions is not None and self._is_gemini_tts_model(model) and not drops_params: + raise UnsupportedParamsError( + message=f"{model} needs instructions to be a string. Send one, or set drop_params to ignore it.", + model=model, + llm_provider="vertex_ai", + ) # Pass through Vertex AI-specific parameters from kwargs if "audioConfig" in kwargs: @@ -311,15 +724,48 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): Google Cloud TTS endpoint: https://texttospeech.googleapis.com/v1/text:synthesize """ + if self._is_gemini_tts_model(model): + return self._get_gemini_tts_url( + model=model, api_base=api_base, litellm_params=_JSON_OBJECT.validate_python(litellm_params) + ) + if api_base: return api_base return self.TTS_API_URL + def _get_gemini_tts_url(self, model: str, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + if api_base: + raise UnsupportedParamsError( + message=( + f"{model} serves mp3, opus, alaw and mulaw through Cloud Text-to-Speech, which does not use " + "this deployment's api_base. Request wav or pcm, or use a deployment without api_base." + ), + model=model, + llm_provider="vertex_ai", + ) + location: Final = litellm_params.get("vertex_location") + vertex_location: Final = location if isinstance(location, str) else self.get_default_vertex_location() + if vertex_location == "global": + return self.TTS_API_URL + cloud_tts_location: Final = _GEMINI_CLOUD_TTS_LOCATIONS.get(vertex_location) + if cloud_tts_location is None: + raise UnsupportedParamsError( + message=( + f"{model} serves mp3, opus, alaw and mulaw through Cloud Text-to-Speech, which has no Gemini TTS " + f"location for vertex_location {vertex_location}. Use a deployment in global, us, eu, " + "northamerica-northeast1, or a US or EU Vertex AI region." + ), + model=model, + llm_provider="vertex_ai", + ) + return f"https://{cloud_tts_location}-texttospeech.googleapis.com/v1/text:synthesize" + def _validate_vertex_input( self, input_data: VertexTextToSpeechInput, optional_params: dict, + detect_ssml: bool, ) -> VertexTextToSpeechInput: """ Validate and transform input for Vertex AI TTS @@ -343,7 +789,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): else: # LiteLLM will auto-detect if text is in ssml format # check if "text" is an ssml - in this case we should pass it as ssml instead of text - if input_data: + if detect_ssml and input_data: _text: Final = input_data.get("text", None) or "" if "" in _text: input_data["ssml"] = input_data.pop("text") @@ -403,8 +849,13 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): headers["x-goog-user-project"] = vertex_project ####### Build the request ################ - vertex_input = VertexTextToSpeechInput(text=input) - vertex_input = self._validate_vertex_input(vertex_input, optional_params) + prompt: Final = _JSON_OBJECT.validate_python(optional_params).get("prompt") + validated_input: Final = self._validate_vertex_input( + VertexTextToSpeechInput(text=input), optional_params, detect_ssml=not self._is_gemini_tts_model(model) + ) + vertex_input: Final[VertexTextToSpeechInput] = ( + {**validated_input, "prompt": prompt} if isinstance(prompt, str) else validated_input + ) # Build voice configuration # Check for voice dict stored in: @@ -412,8 +863,8 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): # 2. optional_params by map_openai_params voice_dict: Final = litellm_params.get("vertex_voice_dict") or optional_params.get("vertex_voice_dict") if voice_dict is not None and isinstance(voice_dict, dict): - vertex_voice = VertexTextToSpeechVoice(**voice_dict) - elif voice is not None and isinstance(voice, str): + vertex_voice = _JSON_OBJECT.validate_python(voice_dict) + elif voice is not None: # Handle string voice (shouldn't normally happen if dispatch was called) parts: Final = voice.split("-") if len(parts) >= 2: @@ -485,8 +936,44 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): content=binary_data, ) - # Initialize the HttpxBinaryResponseContent instance - return HttpxBinaryResponseContent(response) + binary_response: Final = HttpxBinaryResponseContent(response) + if self._is_gemini_tts_model(model): + from litellm.types.utils import CompletionTokensDetailsWrapper, Usage + + request: Final = _GeminiCloudTtsRequest.model_validate_json(raw_response.request.content) + input_data: Final = request.input + audio_config: Final = request.audioConfig + container_duration: Final = calculate_request_duration(binary_data) + sample_rate: Final = audio_config.sampleRateHertz or 24000 + duration: Final = ( + container_duration + if container_duration is not None + else _fallback_gemini_tts_audio_duration(binary_data, audio_config.audioEncoding, sample_rate) + ) + if duration is None: + verbose_logger.warning( + "Could not measure Gemini TTS audio duration for %s, so its audio tokens are not billed", model + ) + return binary_response + input_text: Final = " ".join( + value for value in (input_data.text, input_data.ssml, input_data.prompt) if value + ) + prompt_tokens: Final = litellm.token_counter(model=model, text=input_text) + audio_tokens: Final = math.ceil(duration * 25) + binary_response.usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=audio_tokens, + total_tokens=prompt_tokens + audio_tokens, + completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=audio_tokens), + ) + return binary_response + + +class VertexAIGeminiCloudTextToSpeechConfig(VertexAITextToSpeechConfig): + def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException: + return super().get_error_class( + error_message=error_message, status_code=400 if status_code == 404 else status_code, headers=headers + ) class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): @@ -588,13 +1075,14 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): }, ) location: Final = self.safe_get_vertex_ai_location(litellm_params) or self.get_default_vertex_location() - base_url: Final = self.get_api_base(api_base=api_base, vertex_location=location).rstrip("/") + base_url: Final = self.get_api_base(api_base=api_base, vertex_location=location) encoded_project: Final = encode_url_path_segment(project, field_name="project") encoded_location: Final = encode_url_path_segment(location, field_name="location") encoded_model: Final = encode_url_path_segment(base_model, field_name="model") - return ( - f"{base_url}/v1/projects/{encoded_project}/locations/{encoded_location}" - f"/publishers/google/models/{encoded_model}:predict" + return join_api_base( + base_url, + f"v1/projects/{encoded_project}/locations/{encoded_location}" + f"/publishers/google/models/{encoded_model}:predict", ) def transform_text_to_speech_request( diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 5c250fc1a7e..440604688bb 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -288,7 +289,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Transform create request for Vertex AI RAG Corpus """ - url: Final = f"{api_base}/ragCorpora" # Base URL for creating RAG corpus + url: Final = join_api_base(api_base, "ragCorpora") # Base URL for creating RAG corpus # Add metadata if provided metadata: Final = vector_store_create_optional_params.get("metadata") diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py index 45cc5320401..3873a4be0b5 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py @@ -1,4 +1,5 @@ import types +from collections.abc import Mapping from typing import Final import litellm @@ -18,16 +19,17 @@ class VertexAIAi21Config(OpenAIGPTConfig): self, max_tokens: int | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 5d748ad6e74..cfd5b5d3980 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -1,4 +1,5 @@ -from typing import Any, Final +from collections.abc import Sequence +from typing import Final from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.pass_through.messages.transformation import ( @@ -29,7 +30,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 7100c9efde3..4be41e26326 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -8,6 +8,7 @@ their respective publisher-specific count-tokens endpoints. from typing import Any, Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -59,16 +60,14 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): publisher: Final = self._get_publisher_for_model(model) # Use custom api_base if provided, otherwise construct default - if api_base: - base_url = api_base - else: - base_url = get_vertex_base_url(vertex_location) + base_url: Final = api_base or get_vertex_base_url(vertex_location) # Construct the count-tokens endpoint # Format: /v1/projects/{project}/locations/{location}/publishers/{publisher}/models/count-tokens:rawPredict - endpoint: Final = ( - f"{base_url}/v1/projects/{project_id}/locations/{vertex_location}/" - f"publishers/{publisher}/models/count-tokens:rawPredict" + endpoint: Final = join_api_base( + base_url, + f"v1/projects/{project_id}/locations/{vertex_location}/" + f"publishers/{publisher}/models/count-tokens:rawPredict", ) return endpoint diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index ddc83bc53a5..96fccf007a5 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -1,4 +1,5 @@ import types +from collections.abc import Mapping from typing import Final, Literal from pydantic import TypeAdapter @@ -70,16 +71,17 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): | None = None, title: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 77508096e4a..0cdd92bb071 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -7,13 +7,16 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S import asyncio import json import threading +from collections import OrderedDict from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse import litellm from litellm._logging import verbose_logger +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.secret_managers.main import get_secret_str from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES, VertexPartnerProvider @@ -69,16 +72,20 @@ class _VertexCredentialsObject(Protocol): def refresh(self, request: object) -> None: ... +_VertexCredentialCacheKey = tuple[str | None, str | None] + + class VertexBase: def __init__(self) -> None: super().__init__() self.access_token: str | None = None self.refresh_token: str | None = None self._credentials: GoogleCredentialsObject | None = None - self._credentials_project_mapping: dict[ - tuple[VERTEX_CREDENTIALS_TYPES | None, str | None], + self._credentials_project_mapping: OrderedDict[ + _VertexCredentialCacheKey, tuple[_VertexCredentialsObject, str | None], - ] = {} + ] = OrderedDict() + self._credentials_cache_lock = threading.RLock() self.project_id: str | None = None self.async_handler: AsyncHTTPHandler | None = None # Per-credential-key asyncio.Lock for single-flight async refresh. @@ -90,10 +97,10 @@ class VertexBase: # each lock; the entry is pruned when the count reaches zero, so the # dict stays bounded even in long-running high-cardinality deployments # without depending on any private asyncio internals. - self._async_refresh_locks: dict[tuple, asyncio.Lock] = {} - self._async_refresh_lock_refcounts: dict[tuple, int] = {} + self._async_refresh_locks: dict[_VertexCredentialCacheKey, asyncio.Lock] = {} + self._async_refresh_lock_refcounts: dict[_VertexCredentialCacheKey, int] = {} # Tracks in-flight background refresh tasks to avoid duplicate refreshes. - self._background_refresh_tasks: dict[tuple, asyncio.Task] = {} + self._background_refresh_tasks: dict[_VertexCredentialCacheKey, asyncio.Task] = {} # Protects the sync get_access_token refresh path. # Use RLock so that the reauthentication retry path (which calls # back into get_access_token while still holding the lock) can @@ -295,23 +302,22 @@ class VertexBase: if api_base is None: api_base = get_vertex_base_url(vertex_location) + location_path: Final = f"projects/{vertex_project}/locations/{vertex_location}" + raw_predict_method: Final = "streamRawPredict" if stream else "rawPredict" if partner == VertexPartnerProvider.llama: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" + return join_api_base(api_base, f"v1/{location_path}/endpoints/openapi/chat/completions") elif partner == VertexPartnerProvider.mistralai: - if stream: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/mistralai/models/{model}:streamRawPredict" - else: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/mistralai/models/{model}:rawPredict" + return join_api_base( + api_base, f"v1/{location_path}/publishers/mistralai/models/{model}:{raw_predict_method}" + ) elif partner == VertexPartnerProvider.ai21: - if stream: - return f"{api_base}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/publishers/ai21/models/{model}:streamRawPredict" - else: - return f"{api_base}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/publishers/ai21/models/{model}:rawPredict" + return join_api_base( + api_base, f"v1beta1/{location_path}/publishers/ai21/models/{model}:{raw_predict_method}" + ) elif partner == VertexPartnerProvider.claude: - if stream: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/anthropic/models/{model}:streamRawPredict" - else: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/anthropic/models/{model}:rawPredict" + return join_api_base( + api_base, f"v1/{location_path}/publishers/anthropic/models/{model}:{raw_predict_method}" + ) def get_complete_vertex_url( self, @@ -378,7 +384,7 @@ class VertexBase: with self._sync_refresh_lock: credentials.refresh(Request()) - def _acquire_async_refresh_lock(self, credential_cache_key: tuple) -> asyncio.Lock: + def _acquire_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey) -> asyncio.Lock: """Increment the refcount and return the lock for ``credential_cache_key``. Every call must be paired with ``_release_async_refresh_lock`` once the @@ -391,7 +397,7 @@ class VertexBase: ) return lock - def _release_async_refresh_lock(self, credential_cache_key: tuple, lock: asyncio.Lock) -> None: + def _release_async_refresh_lock(self, credential_cache_key: _VertexCredentialCacheKey, lock: asyncio.Lock) -> None: """Decrement the refcount and drop the lock entry when it reaches zero. Must be called only after the caller has released ``lock`` (i.e. once @@ -409,7 +415,7 @@ class VertexBase: def _try_get_cached_token( self, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, project_id: str | None, ) -> tuple[str, str] | None: """ @@ -432,7 +438,7 @@ class VertexBase: def _try_get_usable_cached_token( self, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, project_id: str | None, ) -> tuple[str, str, "TokenState", _VertexCredentialsObject, str | None] | None: """ @@ -458,19 +464,37 @@ class VertexBase: return creds.token, resolved_project, token_state, creds, cached_project_id def _unpack_cached_credentials( - self, credential_cache_key: tuple + self, credential_cache_key: _VertexCredentialCacheKey ) -> tuple[_VertexCredentialsObject | None, str | None]: """ Return (credentials, project_id) from the cache, or (None, None) if not cached. Handles both tuple and legacy cache formats. """ - if credential_cache_key not in self._credentials_project_mapping: - return None, None - cached_entry: Final = self._credentials_project_mapping[credential_cache_key] + with self._credentials_cache_lock: + cached_entry: Final = self._credentials_project_mapping.get(credential_cache_key) + if cached_entry is None: + return None, None + self._credentials_project_mapping.move_to_end(credential_cache_key) if isinstance(cached_entry, tuple): return cached_entry return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None) + def _cache_credentials( + self, + credential_cache_key: _VertexCredentialCacheKey, + credentials: _VertexCredentialsObject, + project_id: str | None, + *, + replace: bool = True, + ) -> None: + with self._credentials_cache_lock: + if not replace and credential_cache_key in self._credentials_project_mapping: + return + self._credentials_project_mapping[credential_cache_key] = (credentials, project_id) + self._credentials_project_mapping.move_to_end(credential_cache_key) + while len(self._credentials_project_mapping) > max(0, DEFAULT_MAX_LRU_CACHE_SIZE): + self._credentials_project_mapping.popitem(last=False) + def _get_token_state(self, credentials: _VertexCredentialsObject) -> "TokenState": """ Return the token state using google-auth's TokenState enum. @@ -494,7 +518,7 @@ class VertexBase: self, credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, ) -> tuple[_VertexCredentialsObject, str | None]: """Load credentials via load_auth (in thread) and cache the result.""" try: @@ -507,17 +531,12 @@ class VertexBase: raise if _credentials is None: raise ValueError("Could not resolve credentials") - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(credential_cache_key, _credentials, credential_project_id) return _credentials, credential_project_id async def _background_refresh_credentials( self, credentials: _VertexCredentialsObject, - credential_cache_key: tuple, - credential_project_id: str | None, ) -> None: """ Refresh credentials in the background without blocking the calling request. @@ -528,25 +547,13 @@ class VertexBase: try: verbose_logger.debug("Background proactive credential refresh") await asyncify(self.refresh_auth)(credentials) - # Only update the cache if it still points at the credentials - # object we just refreshed. The per-key async lock is not held - # here, so a concurrent INVALID path may have already replaced - # this entry (e.g. via _handle_reauthentication_async, which - # creates a fresh credentials object). In that case our write - # would clobber the newer entry with a stale reference. - cached_creds, _ = self._unpack_cached_credentials(credential_cache_key) - if cached_creds is credentials: - self._credentials_project_mapping[credential_cache_key] = ( - credentials, - credential_project_id, - ) except Exception: verbose_logger.debug( "Background credential refresh failed, will retry on next request", exc_info=True, ) - async def _await_in_flight_background_refresh(self, credential_cache_key: tuple) -> None: + async def _await_in_flight_background_refresh(self, credential_cache_key: _VertexCredentialCacheKey) -> None: """Wait for an in-flight background refresh to finish, if any. google-auth's ``Credentials.refresh()`` is not safe to invoke @@ -568,8 +575,7 @@ class VertexBase: def _schedule_background_refresh( self, credentials: _VertexCredentialsObject, - credential_cache_key: tuple, - credential_project_id: str | None, + credential_cache_key: _VertexCredentialCacheKey, ) -> None: """Kick off a single background refresh for ``credential_cache_key``. @@ -581,9 +587,7 @@ class VertexBase: if existing is not None and not existing.done(): return self._background_refresh_tasks.pop(credential_cache_key, None) - task: Final = asyncio.create_task( - self._background_refresh_credentials(credentials, credential_cache_key, credential_project_id) - ) + task: Final = asyncio.create_task(self._background_refresh_credentials(credentials)) def _drop_background_refresh_task(_fut: asyncio.Future[None]) -> None: if self._background_refresh_tasks.get(credential_cache_key) is _fut: @@ -658,7 +662,7 @@ class VertexBase: # For Gemini (Google AI Studio), construct the full path like other providers if model is None: raise ValueError("Model parameter is required for Gemini custom API base URLs") - url = f"{api_base}/models/{model}:{endpoint}" + url = join_api_base(api_base, f"models/{model}:{endpoint}") if gemini_api_key is None: raise ValueError( "Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable." @@ -678,20 +682,17 @@ class VertexBase: model_for_url: Final = get_vertex_base_model_name(model=model) # Format: {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint} version: Final = vertex_api_version or "v1" - url = "{}/{}/projects/{}/locations/{}/endpoints/{}:{}".format( - api_base.rstrip("/"), - version, - vertex_project, - vertex_location, - model_for_url, - endpoint, + url = join_api_base( + api_base, + f"{version}/projects/{vertex_project}/locations/{vertex_location}" + f"/endpoints/{model_for_url}:{endpoint}", ) elif urlparse(api_base).path in ("", "/"): - url = api_base.rstrip("/") + urlparse(url).path + url = join_api_base(api_base, urlparse(url).path) elif urlparse(api_base).path.rstrip("/") in ("/v1", "/v1beta1") and "/projects/" in urlparse(url).path: url = graft_default_vertex_path(api_base=api_base, default_url=url) else: - url = f"{api_base}:{endpoint}" + url = f"{api_base.rstrip('/')}:{endpoint}" if stream is True: parsed_stream_url: Final = urlparse(url) stream_query: Final = f"{parsed_stream_url.query}&alt=sse" if parsed_stream_url.query else "alt=sse" @@ -799,7 +800,7 @@ class VertexBase: self, credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, error: Exception, ) -> tuple[str, str]: """ @@ -825,8 +826,8 @@ class VertexBase: ) # Clear the cached credentials - if credential_cache_key in self._credentials_project_mapping: - del self._credentials_project_mapping[credential_cache_key] + with self._credentials_cache_lock: + self._credentials_project_mapping.pop(credential_cache_key, None) # Retry once with _retry_reauth=True to prevent infinite recursion try: @@ -849,7 +850,7 @@ class VertexBase: self, credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, - credential_cache_key: tuple, + credential_cache_key: _VertexCredentialCacheKey, error: Exception, ) -> tuple[str, str]: """ @@ -859,7 +860,8 @@ class VertexBase: "Handling async reauthentication for project_id: %s. Clearing cache and retrying once.", project_id ) - self._credentials_project_mapping.pop(credential_cache_key, None) + with self._credentials_cache_lock: + self._credentials_project_mapping.pop(credential_cache_key, None) try: ( @@ -878,10 +880,7 @@ class VertexBase: # references the OLD credentials object we just replaced, and # leaving it would force the next request to do a redundant # refresh/reauth before realizing the cached creds are stale. - self._credentials_project_mapping[resolved_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(resolved_cache_key, _credentials, credential_project_id) if _credentials.token is None or not isinstance(_credentials.token, str): raise ValueError( @@ -927,21 +926,11 @@ class VertexBase: # Convert dict credentials to string for caching cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key: Final = (cache_credentials, project_id) - _credentials: _VertexCredentialsObject | None = None - verbose_logger.debug("Checking cached credentials for project_id: %s", project_id) - if credential_cache_key in self._credentials_project_mapping: + _credentials, credential_project_id = self._unpack_cached_credentials(credential_cache_key) + if _credentials is not None: verbose_logger.debug("Cached credentials found for project_id: %s.", project_id) - # Retrieve both credentials and cached project_id - cached_entry: Final = self._credentials_project_mapping[credential_cache_key] - verbose_logger.debug("cached_entry: %s", cached_entry) - if isinstance(cached_entry, tuple): - _credentials, credential_project_id = cached_entry - else: - # Backward compatibility with old cache format - _credentials = cached_entry - credential_project_id = _credentials.quota_project_id or getattr(_credentials, "project_id", None) verbose_logger.debug( "Using cached credentials for project_id: %s", credential_project_id, @@ -966,10 +955,7 @@ class VertexBase: f"Could not resolve credentials - either dynamically or from environment, for project_id: {project_id}" ) # Cache the project_id and credentials from load_auth result (resolved project_id) - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(credential_cache_key, _credentials, credential_project_id) ## VALIDATE CREDENTIALS verbose_logger.debug("Validating credentials") @@ -977,11 +963,7 @@ class VertexBase: project_id = credential_project_id # Update cache with resolved project_id for future lookups resolved_cache_key: Final = (cache_credentials, project_id) - if resolved_cache_key not in self._credentials_project_mapping: - self._credentials_project_mapping[resolved_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(resolved_cache_key, _credentials, credential_project_id, replace=False) # Check if credentials are None before accessing attributes if _credentials is None: @@ -994,10 +976,6 @@ class VertexBase: try: verbose_logger.debug("Credentials expired, refreshing") self.refresh_auth(_credentials) - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) except Exception as e: # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login` # in this case, we should try to reload the credentials by clearing the cache and retrying @@ -1046,9 +1024,9 @@ class VertexBase: # callers on the lock just to schedule that refresh. usable: Final = self._try_get_usable_cached_token(credential_cache_key, project_id) if usable is not None: - cached_token, resolved_project, token_state, creds, cached_project_id = usable + cached_token, resolved_project, token_state, creds, _cached_project_id = usable if token_state == TokenState.STALE: - self._schedule_background_refresh(creds, credential_cache_key, cached_project_id) + self._schedule_background_refresh(creds, credential_cache_key) return cached_token, resolved_project # === SLOW PATH (per-key lock) === @@ -1078,10 +1056,7 @@ class VertexBase: # reauth that only repopulated the unresolved key), which # would force the next request through an unnecessary # refresh/reauth cycle. - self._credentials_project_mapping[resolved_cache_key] = ( - _credentials, - credential_project_id, - ) + self._cache_credentials(resolved_cache_key, _credentials, credential_project_id) # Use google-auth's token_state to decide refresh strategy: # - STALE: token is usable but within REFRESH_THRESHOLD (3:45) of @@ -1101,7 +1076,6 @@ class VertexBase: self._schedule_background_refresh( _credentials, credential_cache_key, - credential_project_id, ) return current_token, project_id @@ -1120,10 +1094,6 @@ class VertexBase: try: verbose_logger.debug("Credentials expired, refreshing") await asyncify(self.refresh_auth)(_credentials) - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) except Exception as e: if "Reauthentication is needed" in str(e): verbose_logger.debug("Reauthentication needed, clearing cache and retrying") diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index de368cea5ac..080899d08b0 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -18,6 +18,7 @@ from typing_extensions import ReadOnly import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.vertex_ai.common_utils import ( convert_vertex_datetime_to_openai_datetime, @@ -325,13 +326,10 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): model_name: Final = model.replace("vertex_ai/", "") # Construct the URL - if api_base: - base_url = api_base.rstrip("/") - else: - base_url = get_vertex_base_url(vertex_location) + base_url: Final = api_base or get_vertex_base_url(vertex_location) - url: Final = ( - f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}" + url: Final = join_api_base( + base_url, f"v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}" ) return url @@ -480,7 +478,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): # Construct the full URL including model ID # URL format: https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT/locations/LOCATION/publishers/google/models/MODEL:fetchPredictOperation # Strip trailing slashes from api_base and append model - url: Final = f"{api_base.rstrip('/')}/{model}:fetchPredictOperation" + url: Final = join_api_base(api_base, f"{model}:fetchPredictOperation") # Request body contains the operation name params: Final = {"operationName": operation_name} @@ -778,7 +776,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): if vertex_params: request_data["parameters"] = vertex_params - edit_url: Final = f"{api_base.rstrip('/')}/{model}:predictLongRunning" + edit_url: Final = join_api_base(api_base, f"{model}:predictLongRunning") return edit_url, request_data, None def transform_video_edit_response( diff --git a/litellm/llms/volcengine/chat/transformation.py b/litellm/llms/volcengine/chat/transformation.py index 52b3242c186..4ccab38e761 100644 --- a/litellm/llms/volcengine/chat/transformation.py +++ b/litellm/llms/volcengine/chat/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig @@ -34,7 +35,7 @@ class VolcEngineChatConfig(OpenAILikeChatConfig): top_p: int | None = None, response_format: dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/volcengine/embedding/transformation.py b/litellm/llms/volcengine/embedding/transformation.py index b298c9996f7..cc1e7d59caa 100644 --- a/litellm/llms/volcengine/embedding/transformation.py +++ b/litellm/llms/volcengine/embedding/transformation.py @@ -10,6 +10,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues @@ -76,12 +77,8 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): Returns: Complete URL for the embedding API endpoint """ - base_url: Final = get_volcengine_base_url(api_base) - # Construct the complete URL with /embeddings endpoint - if base_url.endswith("/api/v3"): - return f"{base_url}/embeddings" - else: - return f"{base_url}/api/v3/embeddings" + versioned_base_url: Final = ensure_api_route(get_volcengine_base_url(api_base), "api/v3") + return join_api_base(versioned_base_url, "embeddings") def map_openai_params( self, diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 6e0aa88b6b6..747ea21e4fb 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -11,7 +11,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( safe_convert_created_field, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -45,6 +45,11 @@ class _EventModelClass(Protocol): def model_validate(self, obj: Mapping[str, object]) -> ResponsesAPIStreamingResponse: ... +class _FieldDefaultCarrier(Protocol): + @property + def default(self) -> object: ... + + class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): _SUPPORTED_OPTIONAL_PARAMS: list[str] = [ # Doc-listed knobs @@ -137,8 +142,8 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): if base_url.endswith("/responses"): return base_url if base_url.endswith("/api/v3"): - return f"{base_url}/responses" - return f"{base_url}/api/v3/responses" + return join_api_base(base_url, "responses") + return join_api_base(base_url, "api/v3/responses") def map_openai_params( self, @@ -345,7 +350,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): order: Literal["asc", "desc"] = "desc", ) -> tuple[str, dict]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") - url: Final = f"{api_base}/{encoded_response_id}/input_items" + url: Final = join_api_base(api_base, f"{encoded_response_id}/input_items") params: Final[dict[str, str | int]] = {} if after is not None: params["after"] = after @@ -380,7 +385,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): headers: dict, ) -> tuple[str, dict]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") - url: Final = f"{api_base}/{encoded_response_id}/cancel" + url: Final = join_api_base(api_base, f"{encoded_response_id}/cancel") data: Final[dict] = {} return url, data @@ -430,6 +435,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): annotation: Final[object] = field.annotation return annotation + @staticmethod + def _field_default(field: _FieldDefaultCarrier) -> object: + return field.default + @staticmethod def _fill_missing_fields(chunk: Mapping[str, object], event_model: object | None) -> Mapping[str, object]: """ @@ -451,7 +460,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): continue # Explicit default or factory - field_default: object = field.default + field_default: object = VolcEngineResponsesAPIConfig._field_default(field) if field_default is not pyd_fields.PydanticUndefined and field_default is not None: patched[name] = field_default continue diff --git a/litellm/llms/voyage/embedding/transformation.py b/litellm/llms/voyage/embedding/transformation.py index 7e8425bb3d0..1ba66e5ce65 100644 --- a/litellm/llms/voyage/embedding/transformation.py +++ b/litellm/llms/voyage/embedding/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_api_key @@ -50,12 +51,7 @@ class VoyageEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - if api_base: - clean_api_base: Final = api_base.rstrip("/") - if not clean_api_base.endswith("/embeddings"): - return f"{clean_api_base}/embeddings" - return clean_api_base - return f"{get_default_base_url(api_key)}/embeddings" + return ensure_api_route(api_base or get_default_base_url(api_key), "embeddings") def get_supported_openai_params(self, model: str) -> list: return [ diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 8e3802ccbc9..62c03139fbe 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -10,6 +10,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_api_key @@ -58,12 +59,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - if api_base: - clean_api_base: Final = api_base.rstrip("/") - if not clean_api_base.endswith("/contextualizedembeddings"): - return f"{clean_api_base}/contextualizedembeddings" - return clean_api_base - return f"{get_default_base_url(api_key)}/contextualizedembeddings" + return ensure_api_route(api_base or get_default_base_url(api_key), "contextualizedembeddings") def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class signature return ["encoding_format", "dimensions"] diff --git a/litellm/llms/voyage/embedding/transformation_multimodal.py b/litellm/llms/voyage/embedding/transformation_multimodal.py index 29fdb9ecfe1..c04a2e1f26b 100644 --- a/litellm/llms/voyage/embedding/transformation_multimodal.py +++ b/litellm/llms/voyage/embedding/transformation_multimodal.py @@ -13,6 +13,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_api_key @@ -59,12 +60,7 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - if api_base: - clean_api_base: Final = api_base.rstrip("/") - if not clean_api_base.endswith("/multimodalembeddings"): - return f"{clean_api_base}/multimodalembeddings" - return clean_api_base - return f"{get_default_base_url(api_key)}/multimodalembeddings" + return ensure_api_route(api_base or get_default_base_url(api_key), "multimodalembeddings") def get_supported_openai_params(self, model: str) -> list: return ["dimensions"] diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index 121c18e82ae..21c6e6a82b9 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_api_key @@ -28,6 +29,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _OPTIONAL_INT: Final[TypeAdapter[int | None]] = TypeAdapter(int | None) _STR: Final = TypeAdapter(str) +_DECODED_JSON: Final = TypeAdapter(object) class VoyageRerankConfig(BaseRerankConfig): @@ -80,9 +82,9 @@ class VoyageRerankConfig(BaseRerankConfig): api_base = api_base.rstrip("/") if not api_base.endswith("/v1/rerank"): if api_base.endswith("/v1"): - api_base = f"{api_base}/rerank" + api_base = join_api_base(api_base, "rerank") else: - api_base = f"{api_base}/v1/rerank" + api_base = join_api_base(api_base, "v1/rerank") return api_base def transform_rerank_request( @@ -111,7 +113,7 @@ class VoyageRerankConfig(BaseRerankConfig): logging_obj.post_call(original_response=raw_response.text) try: - _json_response: Final = raw_response.json() + _json_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raise VoyageError( message=f"Failed to parse response: {raw_response.text}", diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 77ef03b8e10..35fda05c028 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -12,6 +12,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.llms.openai import ( AllMessageValues, OpenAIAudioTranscriptionOptionalParams, @@ -28,6 +29,7 @@ from ...openai.transcriptions.whisper_transformation import ( from ..common_utils import IBMWatsonXMixin _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTranscriptionConfig): @@ -150,10 +152,9 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran """ # Get base URL url = self._get_base_url(api_base=api_base) - url = url.rstrip("/") # Add the audio transcription endpoint - url = f"{url}/ml/v1/audio/transcriptions" + url = join_api_base(url, "ml/v1/audio/transcriptions") # Add version parameter (only version in query string, not project_id) api_version: Final = optional_params.get("api_version", None) or litellm.WATSONX_DEFAULT_API_VERSION @@ -172,7 +173,7 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran removed before creating the TranscriptionResponse object. """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise ValueError(f"Error transforming response to json: {e}\nResponse: {raw_response.text}") diff --git a/litellm/llms/watsonx/chat/transformation.py b/litellm/llms/watsonx/chat/transformation.py index 4ff2212711d..071ff7108b9 100644 --- a/litellm/llms/watsonx/chat/transformation.py +++ b/litellm/llms/watsonx/chat/transformation.py @@ -7,6 +7,7 @@ Docs: https://cloud.ibm.com/apidocs/watsonx-ai#text-chat from typing import Final from litellm import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.watsonx import ( WatsonXAIEndpoint, @@ -96,7 +97,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): endpoint = endpoint.format(deployment_id=deployment_id) else: endpoint = WatsonXAIEndpoint.CHAT_STREAM.value if stream else WatsonXAIEndpoint.CHAT.value - url = url.rstrip("/") + endpoint + url = join_api_base(url, endpoint) ## add api version url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 5c0416bd6ca..9ea062ea2f8 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock from litellm.types.llms.watsonx import WatsonXAIEndpoint @@ -342,7 +343,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): endpoint = endpoint.format(deployment_id=deployment_id) else: endpoint = WatsonXAIEndpoint.TEXT_GENERATION_STREAM if stream else WatsonXAIEndpoint.TEXT_GENERATION - url = url.rstrip("/") + endpoint + url = join_api_base(url, endpoint) ## add api version url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) diff --git a/litellm/llms/watsonx/embed/transformation.py b/litellm/llms/watsonx/embed/transformation.py index 1b870c7f639..6ec3b50bbe9 100644 --- a/litellm/llms/watsonx/embed/transformation.py +++ b/litellm/llms/watsonx/embed/transformation.py @@ -9,6 +9,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.embedding.transformation import ( BaseEmbeddingConfig, LiteLLMLoggingObj, @@ -92,7 +93,7 @@ class IBMWatsonXEmbeddingConfig(IBMWatsonXMixin, BaseEmbeddingConfig): if model.startswith("deployment/"): deployment_id: Final = "/".join(model.split("/")[1:]) endpoint = endpoint.format(deployment_id=deployment_id) - url = url.rstrip("/") + endpoint + url = join_api_base(url, endpoint) ## add api version url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 5f337603763..87f30483196 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -11,6 +11,7 @@ from typing import Final, cast import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -29,6 +30,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _OPTIONAL_INT: Final[TypeAdapter[int | None]] = TypeAdapter(int | None) _STR: Final = TypeAdapter(str) +_DECODED_JSON: Final = TypeAdapter(object) class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): @@ -45,7 +47,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): base_url: Final = self._get_base_url(api_base=api_base) endpoint: Final = WatsonXAIEndpoint.RERANK.value - url: Final = base_url.rstrip("/") + endpoint + url: Final = join_api_base(base_url, endpoint) params: Final = optional_params or {} @@ -169,7 +171,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): Transform IBM watsonx.ai rerank response to LiteLLM RerankResponse format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Failed to parse response: {e}", diff --git a/litellm/llms/xai/audio_transcription/transformation.py b/litellm/llms/xai/audio_transcription/transformation.py index e34034b5ae9..e29c3d9334a 100644 --- a/litellm/llms/xai/audio_transcription/transformation.py +++ b/litellm/llms/xai/audio_transcription/transformation.py @@ -10,6 +10,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -185,7 +186,7 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> str: base: Final = (XAIModelInfo.get_api_base(api_base) or "").rstrip("/") normalized: Final = base.removesuffix("/v1") - return f"{normalized}/v1/stt" + return join_api_base(normalized, "v1/stt") def validate_environment( self, diff --git a/litellm/llms/xai/batches/transformation.py b/litellm/llms/xai/batches/transformation.py index fdd9939526e..1e803057e3d 100644 --- a/litellm/llms/xai/batches/transformation.py +++ b/litellm/llms/xai/batches/transformation.py @@ -20,7 +20,7 @@ from pydantic import ConfigDict from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import XAI_API_BASE -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.xai.common_utils import XAIModelInfo from litellm.secret_managers.main import get_secret_str @@ -76,10 +76,11 @@ def get_xai_auth_headers( def xai_batches_url(api_base: str | None, batch_id: str | None = None, suffix: str = "") -> str: - base: Final = f"{get_xai_api_base(api_base)}/v1/batches" + root: Final = get_xai_api_base(api_base) if batch_id is None: - return base - return f"{base}/{encode_url_path_segment(batch_id, field_name='batch_id')}{suffix}" + return join_api_base(root, "v1/batches") + encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") + return join_api_base(root, f"v1/batches/{encoded_batch_id}{suffix}") def is_xai_batch_results_id(file_id: str) -> bool: diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 984b62eeb59..2b943fa4107 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -32,6 +32,7 @@ from ...openai.chat.gpt_transformation import ( ) _RESPONSE_BODY: Final = TypeAdapter(dict[object, object], config=ConfigDict(hide_input_in_errors=True)) +_FILTERED_TOOL: Final = TypeAdapter(object) def _usage_restated_from_xai_ticks(usage: Usage | None) -> Usage | None: @@ -204,11 +205,13 @@ class XAIChatConfig(OpenAIGPTConfig): if param == "max_completion_tokens": optional_params["max_tokens"] = value elif param == "tools" and value is not None: - tools = [] - for tool in value: - tool = filter_value_from_dict(tool, "strict") - if tool is not None: - tools.append(tool) + tools = [ + filtered_tool + for filtered_tool in ( + _FILTERED_TOOL.validate_python(filter_value_from_dict(tool, "strict")) for tool in value + ) + if filtered_tool is not None + ] if len(tools) > 0: optional_params["tools"] = tools elif param in supported_openai_params: diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index 4843caeac59..06e3f32a767 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -91,7 +92,7 @@ class XAIModelInfo(BaseLLMModelInfo): "XAI API base or key is not set. Set XAI_API_BASE and provide an xAI API key via api_key, litellm.xai_key, or XAI_API_KEY." ) response: Final = litellm.module_level_client.get( - url=f"{api_base}/v1/models", + url=join_api_base(api_base, "v1/models"), headers={"Authorization": f"Bearer {api_key}"}, ) diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index de2cbf5f9e3..aec535c2086 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -120,12 +120,11 @@ def _web_search_cost_per_call_from_model_info(model_info: "ModelInfo") -> float: search_costs: Final = model_info.get("search_context_cost_per_query") if not isinstance(search_costs, Mapping): return _DEFAULT_WEB_SEARCH_COST_PER_CALL - for key in ( - "search_context_size_medium", - "search_context_size_low", - "search_context_size_high", + for value in ( + search_costs.get("search_context_size_medium"), + search_costs.get("search_context_size_low"), + search_costs.get("search_context_size_high"), ): - value = search_costs.get(key) if value is None: continue try: diff --git a/litellm/llms/xai/files/transformation.py b/litellm/llms/xai/files/transformation.py index 33baca79005..88f2fc51983 100644 --- a/litellm/llms/xai/files/transformation.py +++ b/litellm/llms/xai/files/transformation.py @@ -14,7 +14,7 @@ from pydantic import ConfigDict from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj from litellm.types.llms.base import LiteLLMBaseModel @@ -101,11 +101,11 @@ class XAIFilesConfig(BaseFilesConfig): litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: - return f"{get_xai_api_base(api_base)}/v1/files" + return join_api_base(get_xai_api_base(api_base), "v1/files") def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" + return join_api_base(_api_base_from(litellm_params), f"v1/files/{encoded_file_id}{suffix}") def get_error_class( self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers @@ -204,7 +204,7 @@ class XAIFilesConfig(BaseFilesConfig): optional_params: Mapping[str, object], litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature - return f"{_api_base_from(litellm_params)}/v1/files", _NO_QUERY_PARAMS + return join_api_base(_api_base_from(litellm_params), "v1/files"), _NO_QUERY_PARAMS def transform_list_files_next_request( self, @@ -215,7 +215,7 @@ class XAIFilesConfig(BaseFilesConfig): page: Final = XAIFileList.model_validate(raw_response.json()) if not page.pagination_token or not page.data: return None - return f"{_api_base_from(litellm_params)}/v1/files", {"pagination_token": page.pagination_token} + return join_api_base(_api_base_from(litellm_params), "v1/files"), {"pagination_token": page.pagination_token} def transform_list_files_response( self, diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 2f638da49c8..b6ea516fb79 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -9,6 +9,7 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.secret_managers.main import get_secret_str @@ -249,10 +250,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): else: api_base = api_base or litellm.api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def transform_response_api_response( self, diff --git a/litellm/llms/you_com/search/transformation.py b/litellm/llms/you_com/search/transformation.py index 8b2df949a93..9228deda26b 100644 --- a/litellm/llms/you_com/search/transformation.py +++ b/litellm/llms/you_com/search/transformation.py @@ -10,6 +10,7 @@ from typing import Final, TypedDict import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -109,7 +110,7 @@ class YouComSearchConfig(BaseSearchConfig): api_base = api_base.rstrip("/") if not api_base.endswith("/v1/search") and not api_base.endswith("/v1/agents/search"): - api_base = f"{api_base}/v1/search" + api_base = join_api_base(api_base, "v1/search") return api_base diff --git a/litellm/main.py b/litellm/main.py index d647cab24e3..b3d62daa4d8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -89,6 +89,7 @@ from litellm.litellm_core_utils.get_litellm_params import ( with_control_options, ) from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_provider_affinity_header, invalid_provider_routing_hints_message, ) from litellm.litellm_core_utils.get_provider_specific_headers import ( @@ -112,6 +113,7 @@ from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm import BaseConfig, BaseImageGenerationConfig from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, @@ -129,6 +131,7 @@ from litellm.llms.bedrock_mantle.chat.claude_transformation import bedrock_mantl from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, http2_enabled from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.vertex_ai.common_utils import ( VertexAIModelRoute, @@ -142,6 +145,7 @@ from litellm.router_utils.reasoning_effort_capability import ( resolve_chat_reasoning, ) from litellm.secret_managers.main import get_secret_bool, get_secret_str +from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper from litellm.types.completion import ( CompletionDispatchContext, _CompletionDispatchResult, @@ -168,7 +172,6 @@ from litellm.utils import ( ProviderConfigManager, TextChoices, TextCompletionResponse, - TextCompletionStreamWrapper, TranscriptionResponse, Usage, add_provider_specific_params_to_optional_params, @@ -230,6 +233,7 @@ from .llms.azure.azure import AzureChatCompletion, check_dynamic_azure_params from .llms.azure.chat.o_series_handler import AzureOpenAIO1ChatCompletion from .llms.azure.completion.handler import AzureTextCompletion from .llms.azure_ai.anthropic.handler import AzureAnthropicChatCompletion +from .llms.azure_ai.anthropic.transformation import azure_anthropic_messages_url from .llms.azure_ai.embed import AzureAIEmbedding from .llms.bedrock.chat import BedrockConverseLLM from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig @@ -673,7 +677,11 @@ async def acompletion( } if custom_llm_provider is None: _supplemental: Final = cast( # cast-ok: kwargs values are request-scoped objects - "dict[str, object]", {k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS if k in kwargs} + "dict[str, object]", + { + **{k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS | {"provider_specific_header"} if k in kwargs}, + "extra_headers": extra_headers, + }, ) _, custom_llm_provider, _, _ = get_llm_provider( model=model, @@ -798,6 +806,11 @@ def _transcription_response_from_cached_dict(cached_response_dict: Mapping[str, return TranscriptionResponse(**cached_response_dict) +def _text_completion_response_from_cached_dict(cached_response_dict: Mapping[str, object]) -> TextCompletionResponse: + fields: Final = {key: value for key, value in cached_response_dict.items() if key != "stream"} + return TextCompletionResponse(stream=bool(cached_response_dict.get("stream", False)), **fields) + + async def _async_streaming(response, model, custom_llm_provider, args): try: print_verbose(f"received response in _async_streaming: {response}") @@ -1666,16 +1679,8 @@ def _complete_azure_ai(ctx: CompletionDispatchContext) -> _CompletionDispatchRes ) api_key = AzureFoundryModelInfo.get_api_key(api_key) - # Ensure the URL ends with /v1/messages for Anthropic if api_base: - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/messages"): - if "/anthropic" in api_base: - parts: Final = api_base.split("/anthropic", 1) - api_base = parts[0] + "/anthropic" - else: - api_base = api_base + "/anthropic" - api_base = api_base + "/v1/messages" + api_base = azure_anthropic_messages_url(api_base) response = azure_anthropic_chat_completions.completion( model=model, @@ -2196,6 +2201,27 @@ def _complete_groq(ctx: CompletionDispatchContext) -> _CompletionDispatchResult: ) +def _complete_opencode_go(ctx: CompletionDispatchContext) -> _CompletionDispatchResult: + return base_llm_http_handler.completion( # pyright: ignore[reportUnknownMemberType] # shared legacy handler has untyped request mappings + model=ctx.model, + stream=ctx.stream, + messages=ctx.messages, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # existing dispatch context carries validated messages + acompletion=ctx.acompletion, + api_base=ctx.api_base, + model_response=ctx.model_response, + optional_params=ctx.optional_params, # pyright: ignore[reportUnknownMemberType] # existing dispatch context carries the mapped provider parameters + litellm_params=ctx.litellm_params, + shared_session=ctx.shared_session, + custom_llm_provider=ctx.custom_llm_provider, + timeout=ctx.timeout, # pyright: ignore[reportArgumentType] # completion normalizes timeout before constructing this context + headers=ctx.headers, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy context retains caller headers + encoding=_get_encoding(), + api_key=ctx.api_key, + logging_obj=ctx.logging, + client=_dispatch_client_http(ctx), + ) + + def _complete_bedrock_mantle( ctx: CompletionDispatchContext, ) -> _CompletionDispatchResult: @@ -2895,8 +2921,8 @@ def _complete_anthropic_text( # Check if we should disable automatic URL suffix appending disable_url_suffix: Final = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if api_base is not None and not disable_url_suffix and not api_base.endswith("/v1/complete"): - api_base += "/v1/complete" + if api_base is not None and not disable_url_suffix: + api_base = ensure_api_route(api_base, "v1/complete") elif disable_url_suffix: verbose_logger.debug("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/complete suffix") @@ -2951,8 +2977,8 @@ def _complete_anthropic(ctx: CompletionDispatchContext) -> _CompletionDispatchRe # Check if we should disable automatic URL suffix appending disable_url_suffix: Final = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if api_base is not None and not disable_url_suffix and not api_base.endswith("/v1/messages"): - api_base += "/v1/messages" + if api_base is not None and not disable_url_suffix: + api_base = ensure_api_route(api_base, "v1/messages") elif disable_url_suffix: verbose_logger.debug("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/messages suffix") @@ -3564,7 +3590,7 @@ def _complete_nadir(ctx: CompletionDispatchContext) -> _CompletionDispatchResult encoding=_get_encoding(), api_key=api_key, logging_obj=ctx.logging, - client=ctx.client, + client=_dispatch_client_http(ctx), ) ctx.logging.post_call(input=ctx.messages, api_key=api_key, original_response=response) @@ -4090,9 +4116,7 @@ def _complete_text_completion_inception( api_base = passed_api_base or get_secret_str("INCEPTION_API_BASE") or "https://api.inceptionlabs.ai/v1" # FIM is served at `/v1/fim/completions`; the OpenAI client appends # `/completions`, so point it at the `/v1/fim` base. - api_base = api_base.rstrip("/") - if not api_base.endswith("/fim"): - api_base += "/fim" + api_base = ensure_api_route(api_base, "fim") # Don't forward the server-managed Inception key to a caller-supplied # api_base; only resolve it for the default/server base, or when the @@ -4203,6 +4227,7 @@ _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) +_OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) _AWS_BEDROCK_CLIENT_CREDENTIAL_KEYS: Final = ( @@ -4747,7 +4772,7 @@ def _complete_petals(ctx: CompletionDispatchContext) -> _CompletionDispatchResul ) if stream is True: ## [BETA] # Fake streaming for petals - resp_string: Final = model_response["choices"][0]["message"]["content"] + resp_string: Final = _OPAQUE_VALUE_ADAPTER.validate_python(model_response["choices"][0]["message"]["content"]) return CustomStreamWrapper( resp_string, model, @@ -5657,7 +5682,10 @@ def completion( if deployment_id is not None: # azure llms model = deployment_id custom_llm_provider = "azure" - _supplemental_provider_params: Final = {k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS if k in kwargs} + _supplemental_provider_params: Final = { + **{k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS | {"provider_specific_header"} if k in kwargs}, + "extra_headers": extra_headers, + } model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, @@ -5977,6 +6005,7 @@ def completion( tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), + chatgpt_auth_file=routing_kwargs.get("chatgpt_auth_file"), gigachat_scope=kwargs.get("gigachat_scope"), gigachat_auth_url=kwargs.get("gigachat_auth_url"), gigachat_access_token=kwargs.get("gigachat_access_token"), @@ -5998,19 +6027,28 @@ def completion( }, ) litellm_params: Final = with_control_options(requested_litellm_params, control_options) - if litellm_params.get("provider_affinity_header") is not None: + affinity_api_base: Final = ( + str(client.base_url) + if isinstance(client, openai.OpenAI | openai.AsyncOpenAI) + else get_openai_credentials(api_base=api_base).api_base + if custom_llm_provider == "openai" + else api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) + if litellm_params.get("provider_affinity_header") is not None or default_affinity_header is not None: try: headers = add_provider_affinity_header( headers=headers or litellm.headers or MappingProxyType({}), litellm_params=MappingProxyType( { - "provider_affinity_header": litellm_params["provider_affinity_header"], + "provider_affinity_header": litellm_params.get("provider_affinity_header"), "litellm_session_id": kwargs.get("litellm_session_id"), "session_id": kwargs.get("session_id"), "metadata": metadata, "litellm_metadata": kwargs.get("litellm_metadata"), } ), + default_header=default_affinity_header, ) except ValueError as affinity_error: raise litellm.BadRequestError( @@ -6189,6 +6227,10 @@ def completion( response = _complete_xai(_dispatch_ctx) elif custom_llm_provider == "groq": response = _complete_groq(_dispatch_ctx) + elif custom_llm_provider == "opencode_go": + response = _complete_opencode_go( # rebind-ok: common dispatcher finalizes the selected provider result + _dispatch_ctx + ) elif custom_llm_provider == "bedrock_mantle": response = _complete_bedrock_mantle(_dispatch_ctx) elif custom_llm_provider == "a2a": @@ -7716,7 +7758,7 @@ async def atext_completion(*args, **kwargs) -> TextCompletionResponse | TextComp init_response: Final = await loop.run_in_executor(None, func_with_context) if isinstance(init_response, dict) or isinstance(init_response, TextCompletionResponse): ## CACHING SCENARIO if isinstance(init_response, dict): - response = TextCompletionResponse(**init_response) + response = _text_completion_response_from_cached_dict(init_response) else: response = init_response elif asyncio.iscoroutine(init_response): @@ -7896,7 +7938,6 @@ def text_completion( tokenizer: Final = get_encoding("p50k_base") ## if it's a 2d list - each element in the list is a text_completion() request if len(prompt) > 0 and isinstance(prompt[0], list): - responses: Final = [None for x in prompt] # init responses def process_prompt(i, individual_prompt): decoded_prompt: Final = tokenizer.decode(individual_prompt) @@ -7913,16 +7954,16 @@ def text_completion( text_completion_response["object"] = "text_completion" text_completion_response["created"] = response.get("created", None) text_completion_response["model"] = response.get("model", None) - return response["choices"][0] + return response.choices[0] with concurrent.futures.ThreadPoolExecutor() as executor: completed_futures: Final = [ executor.submit(process_prompt, i, individual_prompt) for i, individual_prompt in enumerate(prompt) ] - for i, future in enumerate(concurrent.futures.as_completed(completed_futures)): - responses[i] = future.result() - text_completion_response.choices = responses + text_completion_response.choices = [ + future.result() for future in concurrent.futures.as_completed(completed_futures) + ] return text_completion_response # else: @@ -8846,8 +8887,11 @@ def speech( generic_optional_params: Final = GenericLiteLLMParams.model_validate(kwargs) - # Handle Gemini models separately (they use speech_to_completion_bridge) - if "gemini" in model: + needs_cloud_text_to_speech: Final = ( + litellm.utils.is_gemini_tts_model(model) + and response_format in VertexAITextToSpeechConfig.GEMINI_FORMAT_MAPPINGS + ) + if "gemini" in model and not needs_cloud_text_to_speech: from .endpoints.speech.speech_to_completion_bridge.handler import ( speech_to_completion_bridge_handler, ) @@ -9104,6 +9148,7 @@ async def ahealth_check( custom_llm_provider=custom_llm_provider_from_params, api_base=api_base_from_params, api_key=api_key_from_params, + litellm_params=GenericLiteLLMParams(**model_params), ) model_params["cache"] = {"no-cache": True} # don't used cached responses for making health check calls diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index cdb9a925403..2e9d6a77a62 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -103,6 +103,7 @@ def _target(request: NativeCall, model: str, provider: str | None) -> Target: _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **detectors.COMMON, + Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.RESPONSE_CACHE: detectors.response_cache, Feature.MODIFY_PARAMS: detectors.modify_params, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3d89af9d691..d9df9d97d7b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -644,9 +644,13 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 3.3e-07, "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2.75e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-nova-2-sonic.md", "supports_audio_input": true, "supports_audio_output": true }, @@ -4339,7 +4343,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -6202,6 +6206,8 @@ "text", "audio" ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6266,6 +6272,8 @@ "text", "audio" ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -9232,7 +9240,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -9679,7 +9687,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5.5e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -12827,6 +12835,7 @@ "supports_vision": true }, "azure_ai/Phi-4": { + "deprecation_date": "2027-01-17", "input_cost_per_token": 1.25e-07, "litellm_provider": "azure_ai", "max_input_tokens": 16384, @@ -12840,6 +12849,7 @@ "supports_vision": false }, "azure_ai/Phi-4-mini-instruct": { + "deprecation_date": "2027-01-17", "input_cost_per_token": 7.5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12851,6 +12861,7 @@ "supports_function_calling": true }, "azure_ai/Phi-4-multimodal-instruct": { + "deprecation_date": "2027-01-17", "input_cost_per_audio_token": 4e-06, "input_cost_per_token": 8e-08, "litellm_provider": "azure_ai", @@ -12865,6 +12876,7 @@ "supports_vision": true }, "azure_ai/Phi-4-mini-reasoning": { + "deprecation_date": "2027-01-17", "input_cost_per_token": 7.5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12876,6 +12888,7 @@ "supports_function_calling": true }, "azure_ai/Phi-4-reasoning": { + "deprecation_date": "2027-01-17", "input_cost_per_token": 1.25e-07, "litellm_provider": "azure_ai", "max_input_tokens": 32768, @@ -15654,6 +15667,28 @@ "/v1/images/generations" ] }, + "cerebras/llama-3.1-70b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "supports_function_calling": true, + "supports_tool_choice": true + }, + "cerebras/llama-3.1-8b": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-07, + "supports_function_calling": true, + "supports_tool_choice": true + }, "cerebras/llama-3.3-70b": { "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", @@ -24892,27 +24927,118 @@ ] }, "elevenlabs/eleven_v3": { - "input_cost_per_character": 0.00018, + "input_cost_per_character": 8e-05, "litellm_provider": "elevenlabs", "metadata": { - "calculation": "$0.18/1000 characters (Scale plan pricing, 1 credit per character)", + "calculation": "$0.08/1000 characters (API usage pricing, billed in USD per character)", "notes": "ElevenLabs Eleven v3 - most expressive TTS model with 70+ languages and audio tags support" }, "mode": "audio_speech", - "source": "https://elevenlabs.io/pricing", + "source": "https://elevenlabs.io/pricing/api", "supported_endpoints": [ "/v1/audio/speech" ] }, "elevenlabs/eleven_multilingual_v2": { - "input_cost_per_character": 0.00018, + "input_cost_per_character": 8e-05, "litellm_provider": "elevenlabs", "metadata": { - "calculation": "$0.18/1000 characters (Scale plan pricing, 1 credit per character)", + "calculation": "$0.08/1000 characters (API usage pricing, billed in USD per character)", "notes": "ElevenLabs Eleven Multilingual v2 - default TTS model with 29 languages support" }, "mode": "audio_speech", - "source": "https://elevenlabs.io/pricing", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_v4": { + "input_cost_per_character": 8e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.08/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven v4 - most emotive TTS model with audio tags and 90+ languages" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_v4_turbo": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven v4 Turbo - low-latency v4 tuned for real-time and agents" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_v3_conversational": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven v3 Conversational - low-latency v3 tuned for realtime conversation" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_flash_v2_5": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven Flash v2.5 - low-latency multilingual TTS model" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_flash_v2": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven Flash v2 - low-latency English-only TTS model" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_turbo_v2_5": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven Turbo v2.5 - multilingual TTS model" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_turbo_v2": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven Turbo v2 - English-only TTS model" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", "supported_endpoints": [ "/v1/audio/speech" ] @@ -31041,11 +31167,15 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_batches": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31370,11 +31500,15 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_audio_token_cost": 3e-08, + "cache_read_input_audio_token_cost_batches": 3e-08, + "cache_read_input_audio_token_cost_flex": 3e-08, + "cache_read_input_audio_token_cost_priority": 5.4e-08, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_audio_token": 3e-07, + "input_cost_per_audio_token_flex": 1.5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31664,12 +31798,16 @@ }, "gemini/gemini-3.1-flash-lite": { "cache_read_input_audio_token_cost": 5e-08, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "cache_read_input_audio_token_cost_priority": 9e-08, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "deprecation_date": "2027-05-07", "input_cost_per_audio_token": 5e-07, + "input_cost_per_audio_token_flex": 2.5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -31790,10 +31928,14 @@ }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 5e-08, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -33100,6 +33242,7 @@ "github_copilot/claude-opus-4.8-fast": { "cache_creation_input_token_cost": 1.25e-05, "cache_read_input_token_cost": 1e-06, + "deprecation_date": "2026-11-16", "input_cost_per_token": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 200000, @@ -33559,7 +33702,7 @@ "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33583,7 +33726,7 @@ "input_cost_per_token": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33607,7 +33750,7 @@ "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33631,7 +33774,7 @@ "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -36699,6 +36842,16 @@ }, "gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_12hr": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, + "cache_creation_input_token_cost_above_12hr_priority": 5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -36784,6 +36937,14 @@ }, "gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, @@ -36861,17 +37022,25 @@ }, "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, - "cache_creation_input_token_cost_decisions": 1e-07, + "cache_creation_input_token_cost_above_12hr": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, + "cache_creation_input_token_cost_above_12hr_priority": 5e-07, + "cache_creation_input_token_cost_decisions": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, - "cache_creation_input_token_cost_above_272k_tokens_decisions": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_decisions": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-07, "cache_creation_input_token_cost_flex": 6.25e-08, "cache_creation_input_token_cost_priority": 2.5e-07, "cache_read_input_token_cost": 1e-08, - "cache_read_input_token_cost_decisions": 1e-07, + "cache_read_input_token_cost_decisions": 1e-08, "cache_read_input_token_cost_above_272k_tokens": 2e-08, - "cache_read_input_token_cost_above_272k_tokens_decisions": 2e-07, + "cache_read_input_token_cost_above_272k_tokens_decisions": 2e-08, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 4e-08, "cache_read_input_token_cost_batches": 5e-09, @@ -36947,6 +37116,12 @@ }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_12hr": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, @@ -37021,6 +37196,14 @@ }, "gpt-5.6-sol": { "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_12hr": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_batches": 5e-06, + "cache_creation_input_token_cost_above_12hr_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, @@ -37097,6 +37280,14 @@ }, "gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, @@ -37172,6 +37363,14 @@ }, "gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_12hr": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_priority": 1e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, @@ -37247,6 +37446,7 @@ }, "gpt-5.6-cyber": { "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_above_12hr": 3.125e-05, "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, @@ -42337,6 +42537,134 @@ ], "supports_audio_output": true }, + "mistral/voxtral-mini-transcribe-3": { + "input_cost_per_second": 5.83333e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-mini-transcribe-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-transcribe-3-0": { + "input_cost_per_second": 5.83333e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-mini-transcribe-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-small-transcribe-3": { + "input_cost_per_second": 6.66667e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-small-transcribe-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-small-transcribe-3-0": { + "input_cost_per_second": 6.66667e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-small-transcribe-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-transcribe-realtime-3": { + "input_cost_per_second": 0.000133333, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-mini-transcribe-realtime-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-transcribe-realtime-3-0": { + "input_cost_per_second": 0.000133333, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-mini-transcribe-realtime-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-tts-3": { + "input_cost_per_character": 1.3e-05, + "litellm_provider": "mistral", + "mode": "audio_speech", + "source": "https://docs.mistral.ai/models/voxtral-mini-tts-3-0", + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_output": true + }, + "mistral/voxtral-mini-tts-3-0": { + "input_cost_per_character": 1.3e-05, + "litellm_provider": "mistral", + "mode": "audio_speech", + "source": "https://docs.mistral.ai/models/voxtral-mini-tts-3-0", + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_output": true + }, "mistral/voxtral-small-2507": { "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 1e-07, @@ -42416,6 +42744,29 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5e-07, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://docs.mistral.ai/models/zai-glm-5-3-flash", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/zai-glm-5": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -45623,6 +45974,7 @@ "output_cost_per_token": 0.0 }, "openai.gpt-oss-120b-1:0": { + "supports_vision": false, "supported_endpoints": [ "/v1/chat/completions" ], @@ -45648,6 +46000,7 @@ "supports_tool_choice": true }, "openai.gpt-oss-20b-1:0": { + "supports_vision": false, "supported_endpoints": [ "/v1/chat/completions" ], @@ -48499,6 +48852,17 @@ "supports_tool_choice": true, "supports_vision": true }, + "together_ai/MiniMaxAI/MiniMax-M3.1-Flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://www.together.ai/pricing", + "supports_prompt_caching": true + }, "together_ai/Prism-ML/Ternary-Bonsai-27B": { "input_cost_per_token": 0.0, "litellm_provider": "together_ai", @@ -49352,6 +49716,7 @@ "supports_vision": false }, "us-gov.openai.gpt-oss-20b-1:0": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, @@ -49371,6 +49736,7 @@ "supports_tool_choice": true }, "us-gov.openai.gpt-oss-120b-1:0": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, @@ -52906,12 +53272,14 @@ }, "vertex_ai/meta/llama-4-maverick-17b-16e-instruct-maas": { "input_cost_per_token": 3.5e-07, + "input_cost_per_token_batches": 1.75e-07, "litellm_provider": "vertex_ai-llama_models", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 1.15e-06, + "output_cost_per_token_batches": 5.75e-07, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", "supported_modalities": [ "text", @@ -52926,12 +53294,14 @@ }, "vertex_ai/meta/llama-4-scout-17b-128e-instruct-maas": { "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, "litellm_provider": "vertex_ai-llama_models", "max_input_tokens": 10000000, "max_output_tokens": 10000000, "max_tokens": 10000000, "mode": "chat", "output_cost_per_token": 7e-07, + "output_cost_per_token_batches": 3.5e-07, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", "supported_modalities": [ "text", @@ -59213,6 +59583,7 @@ "supports_audio_input": true, "supports_audio_output": true, "gemini_native_audio": true, + "input_cost_per_video_token": 3e-06, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -59408,6 +59779,7 @@ "tpm": 250000, "rpm": 10, "gemini_native_audio": true, + "input_cost_per_video_token": 3e-06, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -59917,6 +60289,7 @@ } }, "bedrock_mantle/openai.gpt-oss-120b": { + "supports_vision": false, "input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, @@ -59942,6 +60315,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-20b": { + "supports_vision": false, "input_cost_per_token": 7e-08, "input_cost_per_token_batches": 3.5e-08, "input_cost_per_token_flex": 3.5e-08, @@ -59967,6 +60341,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-120b": { + "supports_vision": false, "input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, @@ -59990,6 +60365,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { + "supports_vision": false, "input_cost_per_token": 7e-08, "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, @@ -64195,13 +64571,17 @@ }, "fireworks_ai/accounts/fireworks/models/muse-spark-open": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_priority": 2.25e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_priority": 1.875e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "mode": "chat", "output_cost_per_token": 4.25e-06, - "source": "https://fireworks.ai/models/fireworks/muse-spark-open", + "output_cost_per_token_priority": 6.375e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_reasoning": true, "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": { @@ -67837,6 +68217,7 @@ "supports_vision": false }, "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -67856,6 +68237,7 @@ "supports_tool_choice": true }, "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -68134,6 +68516,7 @@ "supports_vision": false }, "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -68153,6 +68536,7 @@ "supports_tool_choice": true }, "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -68541,6 +68925,7 @@ "supports_vision": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", @@ -68560,6 +68945,7 @@ "supports_tool_choice": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", @@ -68700,6 +69086,7 @@ "supports_vision": true }, "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", @@ -68719,6 +69106,7 @@ "supports_tool_choice": true }, "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", @@ -73326,6 +73714,7 @@ ] }, "vertex_ai/veo-2.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -74673,6 +75062,16 @@ }, "gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, @@ -75038,10 +75437,13 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 3.3e-07, "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2.75e-06, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-nova-2-5-sonic.md", "supports_audio_input": true, "supports_audio_output": true }, @@ -76059,9 +76461,10 @@ "azure_ai/Microsoft-Decision-1": { "input_cost_per_token": 4.2e-08, "litellm_provider": "azure_ai", + "max_input_tokens": 32768, "mode": "evaluation", "output_cost_per_token": 0.0, - "source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'" + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" }, "gpt-rosalind-discovery": { "cache_read_input_token_cost": 5e-07, @@ -76709,14 +77112,14 @@ "supports_response_schema": true }, "baseten/reflection/Beam-1.0": { - "cache_read_input_token_cost": 1.3e-07, - "input_cost_per_token": 1.3e-06, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, "litellm_provider": "baseten", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4.3e-06, + "output_cost_per_token": 1.2e-06, "source": "https://inference.baseten.co/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -77018,5 +77421,1481 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false + }, + "azure_ai/MAI-Code-1.1-Flash": { + "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2026-12-09", + "input_cost_per_token": 2e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 256000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "opencode_go/grok-4.7": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 5e-07, + "max_tokens": 500000, + "max_output_tokens": 500000, + "max_input_tokens": 500000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 200000 + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 5e-07 + }, + { + "range": [ + 200000, + 500000 + ], + "input_cost_per_token": 4e-06, + "output_cost_per_token": 1.2e-05, + "cache_read_input_token_cost": 1e-06 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/gpt-6-luna": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_creation_input_token_cost": 1.25e-07, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 922000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 272000 + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_creation_input_token_cost": 1.25e-07 + }, + { + "range": [ + 272000, + 1050000 + ], + "input_cost_per_token": 2.0000000000000002e-07, + "output_cost_per_token": 7.5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_creation_input_token_cost": 2.5e-07 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/gpt-5.6-luna": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 2.0000000000000002e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 2e-08, + "cache_creation_input_token_cost": 2.5e-07, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 922000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 272000 + ], + "input_cost_per_token": 2.0000000000000002e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 2e-08, + "cache_creation_input_token_cost": 2.5e-07 + }, + { + "range": [ + 272000, + 1050000 + ], + "input_cost_per_token": 4.0000000000000003e-07, + "output_cost_per_token": 1.8000000000000001e-06, + "cache_read_input_token_cost": 4e-08, + "cache_creation_input_token_cost": 5e-07 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/claude-haiku-5-5": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_creation_input_token_cost": 1.25e-07, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 100000 + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_creation_input_token_cost": 1.25e-07 + }, + { + "range": [ + 100000, + 1000000 + ], + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 5.0000000000000004e-08, + "cache_creation_input_token_cost": 6.25e-07 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/gemini-4-argon": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1beta/models" + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1e-07, + "max_tokens": 65536, + "max_output_tokens": 65536, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/glm-5.3-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 3e-08, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/glm-5.3": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/glm-5.2": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/kimi-k3.1-pro": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2.0000000000000002e-07, + "cache_creation_input_token_cost": 2e-06, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/kimi-k3": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3e-06, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/kimi-k2.7-code": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 9.499999999999999e-07, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 1.9e-07, + "max_tokens": 262144, + "max_output_tokens": 262144, + "max_input_tokens": 262144, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/longcat-2.0": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/deepseek-v4.1-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-09, + "max_tokens": 384000, + "max_output_tokens": 384000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/deepseek-v4-pro": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.32e-06, + "output_cost_per_token": 3.96e-06, + "cache_read_input_token_cost": 4.4e-08, + "max_tokens": 384000, + "max_output_tokens": 384000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/mimo-v2.6-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 2.8e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/mimo-v2.6-pro": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 4.35e-07, + "output_cost_per_token": 8.7e-07, + "cache_read_input_token_cost": 3.625e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/mimo-v2.5": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 2.8e-09, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/mimo-v2.5-pro": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 4.35e-07, + "output_cost_per_token": 8.7e-07, + "cache_read_input_token_cost": 3.625e-09, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/minimax-m3.1-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 1.5e-08, + "max_tokens": 512000, + "max_output_tokens": 512000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/minimax-m3": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-08, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/muse-spark-1.3-contributor": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 2.0000000000000002e-07, + "cache_read_input_token_cost": 2e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/muse-spark-1.2-contributor": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 2.0000000000000002e-07, + "cache_read_input_token_cost": 2e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/qwen3.8-max": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost": 2.5e-06, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/qwen3.8-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 4.6999999999999995e-07, + "cache_read_input_token_cost": 1.6e-08, + "cache_creation_input_token_cost": 2.0000000000000002e-07, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/qwen3.7-plus": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 4.0000000000000003e-07, + "output_cost_per_token": 1.6000000000000001e-06, + "cache_read_input_token_cost": 4e-08, + "cache_creation_input_token_cost": 5e-07, + "max_tokens": 65536, + "max_output_tokens": 65536, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 256000 + ], + "input_cost_per_token": 4.0000000000000003e-07, + "output_cost_per_token": 1.6000000000000001e-06, + "cache_read_input_token_cost": 4e-08, + "cache_creation_input_token_cost": 5e-07 + }, + { + "range": [ + 256000, + 1000000 + ], + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 4.8e-06, + "cache_read_input_token_cost": 1.2e-07, + "cache_creation_input_token_cost": 1.5e-06 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/hy4-preview": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 8.34e-07, + "output_cost_per_token": 2.501e-06, + "cache_read_input_token_cost": 4.2000000000000006e-08, + "max_tokens": 64000, + "max_output_tokens": 64000, + "max_input_tokens": 1024000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/hy3": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 5.8e-07, + "cache_read_input_token_cost": 3.5e-08, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 192000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/longcat-2.5-preview-free": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "cache_read_input_token_cost": 0, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/step-5-preview": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 1.0000000000000001e-07, + "max_tokens": 65536, + "max_output_tokens": 65536, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/beam-free": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "cache_read_input_token_cost": 0, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "bedrock/claude_platform/claude-fable-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-fable-5-1": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_computer_use": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-5-5": { + "input_cost_per_token_above_100k_tokens_batches": 2.5e-07, + "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, + "cache_read_input_token_cost_above_100k_tokens_batches": 2.5e-08, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_1hr": 2e-07, + "cache_creation_input_token_cost_batches": 6.25e-08, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_batches": 5e-09, + "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_batches": 2.5e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": true, + "thinking_always_on": false, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true, + "input_cost_per_token_above_100k_tokens": 5e-07, + "output_cost_per_token_above_100k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, + "cache_read_input_token_cost_above_100k_tokens": 5e-08 + }, + "bedrock/claude_platform/claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-7": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 2048, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-8": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_creation_input_token_cost_batches": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost_batches": 1.5e-07, + "deprecation_date": "2026-11-30", + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "provider_specific_entry": { + "us": 1.1 + }, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5-5": { + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": false, + "thinking_always_on": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true + }, + "gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_12hr": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-07, + "cache_creation_input_token_cost_priority": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "cache_read_input_token_cost_above_272k_tokens_batches": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-07, + "cache_read_input_token_cost_batches": 2.5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, + "cache_read_input_token_cost_priority": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "input_cost_per_token_above_272k_tokens_batches": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_priority": 4e-06, + "input_cost_per_token_batches": 5e-07, + "input_cost_per_token_flex": 5e-07, + "input_cost_per_token_priority": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "output_cost_per_token_above_272k_tokens_batches": 3.75e-06, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.5e-05, + "output_cost_per_token_batches": 2.5e-06, + "output_cost_per_token_flex": 2.5e-06, + "output_cost_per_token_priority": 1e-05, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6.1-terra": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "global.openai.gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_bedrock_converse_tool_result_images": false, + "supports_bedrock_runtime_chat_completions_response_format": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "us.openai.gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, + "cache_read_input_token_cost": 5.5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "input_cost_per_token_above_272k_tokens": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "output_cost_per_token_above_272k_tokens": 8.25e-06, + "source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_bedrock_converse_tool_result_images": false, + "supports_bedrock_runtime_chat_completions_response_format": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "baseten/decision-models/baseten-elephant-9b-latest": { + "input_cost_per_token": 0.0, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true + }, + "baseten/decision-models/baseten-mammoth-27b-latest": { + "input_cost_per_token": 0.0, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true + }, + "github_copilot/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true } } diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py index 435db632026..b2bfb572b7d 100644 --- a/litellm/models/managed_files.py +++ b/litellm/models/managed_files.py @@ -6,7 +6,7 @@ Canonical definitions for the ``litellm_managed*`` tables. Re-exported from """ from datetime import datetime -from typing import Any, Literal +from typing import Literal from litellm.types.llms.base import LiteLLMPydanticObjectBase from litellm.types.llms.openai import OpenAIFileObject, ResponsesAPIResponse @@ -39,7 +39,7 @@ class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): """Table for managing vector stores with target_model_names support.""" unified_resource_id: str - resource_object: Any | None = None + resource_object: object | None = None model_mappings: dict[str, str] flat_model_resource_ids: list[str] created_by: str | None = None @@ -54,11 +54,11 @@ class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): custom_llm_provider: str vector_store_name: str | None = None vector_store_description: str | None = None - vector_store_metadata: dict[str, Any] | None = None + vector_store_metadata: dict[str, object] | None = None created_at: datetime | None = None updated_at: datetime | None = None litellm_credential_name: str | None = None - litellm_params: dict[str, Any] | None = None + litellm_params: dict[str, object] | None = None team_id: str | None = None user_id: str | None = None is_config: bool = False diff --git a/litellm/models/organization_membership.py b/litellm/models/organization_membership.py index 6a0d22cb63b..f3abbc2e93d 100644 --- a/litellm/models/organization_membership.py +++ b/litellm/models/organization_membership.py @@ -6,6 +6,7 @@ Canonical definition for ``litellm_organizationmembership``. Re-exported from """ from datetime import datetime +from typing import Final from pydantic import ConfigDict, model_validator @@ -31,9 +32,10 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): @model_validator(mode="after") def populate_user_email(self) -> "LiteLLM_OrganizationMembershipTable": - if self.user_email is None and self.user is not None: - if isinstance(self.user, dict): - self.user_email = self.user.get("user_email") + user: Final[object] = self.user + if self.user_email is None and user is not None: + if isinstance(user, dict): + self.user_email = user.get("user_email") else: - self.user_email = getattr(self.user, "user_email", None) + self.user_email = getattr(user, "user_email", None) return self diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 53b56c7b31d..f883031a529 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -27,6 +27,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_passthrough_resource_metadata_url, get_passthrough_www_authenticate, get_request_base_url, + get_route_relative_request_path, well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( @@ -99,6 +100,7 @@ if TYPE_CHECKING: _EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({}) OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) +_CACHED_VALUE: Final = TypeAdapter(object) def _normalize_caller_admission_credential(value: str) -> str: @@ -386,7 +388,7 @@ def _gateway_dcr_challenge( """The RFC 9728 challenge pointing the client at the protected-resource metadata matching the scope it requested: the per-server document (same URL spelling the request arrived on) when the single target advertises gateway-owned sign-in, - else the gateway's aggregate document. Either way the client discovers the gateway + else a gateway document containing the public resource. Either way the client discovers the gateway as its authorization server and starts the same sign-in flow. ``invalid_token`` adds the RFC 6750 error code for a request that DID @@ -396,10 +398,13 @@ def _gateway_dcr_challenge( client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) target: Final = _gateway_dcr_challenge_target(route, mcp_servers, client_ip) server_scope: Final = gateway_dcr_challenge_scope(route, mcp_servers, client_ip) if target is None else None + public_path: Final = get_route_relative_request_path(request.scope).rstrip("/") + aggregate_resource: Final = public_path == "/mcp" or public_path.startswith("/mcp/") + metadata_suffix: Final = f"{well_known_root_suffix()}/mcp" if aggregate_resource else "" resource_metadata_url: Final = ( get_passthrough_resource_metadata_url(request.scope, target) if target is not None - else f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp" + else f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{metadata_suffix}" ) error_attr: Final = 'error="invalid_token", ' if invalid_token else "" scope_attr: Final = f', scope="{server_scope}"' if server_scope is not None else "" @@ -1368,7 +1373,7 @@ class MCPRequestHandler: if auth is None: return None - refreshed = auth.model_copy() + refreshed: Final = auth.model_copy(update={"requires_fresh_policy": True}) if (auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id: # A session credential (``lite login`` or a redeemed ID-JAG) has no key row, its authority is # sealed in the token: refresh the live user policy and keep the token's team, routes and models. @@ -1386,7 +1391,12 @@ class MCPRequestHandler: refreshed.project_id = current.project_id refreshed.user_id = current.user_id refreshed.access_group_ids = current.access_group_ids - refreshed.requires_fresh_policy = True + else: + return refreshed + if refreshed.mcp_toolset_id is not None: + from litellm.proxy._experimental.mcp_server.server import apply_toolset_scope + + return await apply_toolset_scope(refreshed, refreshed.mcp_toolset_id) return refreshed @staticmethod @@ -4253,7 +4263,7 @@ class MCPRequestHandler: cache_key: Final = f"agent_object_permission_id:{agent_id}" try: - cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final = _CACHED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key)) if cached == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: return None if isinstance(cached, str) and cached: diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 71eb2112bd5..2db548f5863 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -1173,12 +1173,9 @@ async def list_gateway_catalog( try: page: Final = await fetch_optional_catalog_page(context, request, server, allowed, cursor) + page_meta: Final[Mapping[str, object]] = page.meta or {} return page.model_copy( - update={ - "meta": { - key: value for key, value in (page.meta or {}).items() if key != SERVER_OUTCOMES_META_KEY - } - } + update={"meta": {key: value for key, value in page_meta.items() if key != SERVER_OUTCOMES_META_KEY}} ) except Exception as error: if cursor is not None: diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 7302b74f15a..0666cc7e6c0 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -165,6 +165,7 @@ _CLIENT_FORWARDED_AUTH_TYPES: Final["frozenset[str]"] = frozenset({"true_passthr _MINTED_TOKEN_CREDENTIAL_FIELDS: Final["frozenset[str]"] = frozenset({"access_token", "refresh_token", "expires_in"}) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) def _bind_submitted_oauth_client( @@ -339,7 +340,7 @@ def _decrypt_env_vars_on_returned_row(row: object) -> None: return if isinstance(env_vars, str): try: - env_vars = json.loads(env_vars) + env_vars = _DECODED_JSON.validate_python(json.loads(env_vars)) except (json.JSONDecodeError, TypeError): return if not isinstance(env_vars, list): @@ -476,8 +477,9 @@ def _prepare_mcp_server_data( # Serialize JSON fields from ``data_dict`` (not ``data``) so the # exclude_unset filter is respected. Reading back from ``data`` would # reintroduce defaults (e.g. ``env={}``) for fields the caller never set. - if data_dict.get("static_headers") is not None: - data_dict["static_headers"] = encrypt_secret_map(data_dict["static_headers"]) + static_headers: Final[Mapping[str, str] | None] = data_dict.get("static_headers") + if static_headers is not None: + data_dict["static_headers"] = encrypt_secret_map(static_headers) # env_vars is read from ``data_dict`` (not ``data``) like every other JSON # column so the exclude_unset filter is respected: a partial update that @@ -489,11 +491,13 @@ def _prepare_mcp_server_data( _encrypt_global_env_var_values(serialized_env_vars) data_dict["env_vars"] = safe_dumps(serialized_env_vars) - if data_dict.get("mcp_info") is not None: - data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"]) + mcp_info: Final[object] = data_dict.get("mcp_info") + if mcp_info is not None: + data_dict["mcp_info"] = safe_dumps(mcp_info) - if data_dict.get("env") is not None: - data_dict["env"] = encrypt_secret_map(data_dict["env"]) + env: Final[Mapping[str, str] | None] = data_dict.get("env") + if env is not None: + data_dict["env"] = encrypt_secret_map(env) if "tool_name_to_display_name" in data_dict: data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {}) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3b2c2e09df5..f5c54451006 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -701,7 +701,7 @@ async def _store_per_user_token_server_side( expires_in = None refresh_token: Final[str | None] = token_response.get("refresh_token") or None - raw_scope: Final[object] = token_response.get("scope", requested_scope) + raw_scope: Final[object] = _DECODED_JSON.validate_python(token_response.get("scope", requested_scope)) scopes: Final[list[str] | None] = raw_scope.split() if isinstance(raw_scope, str) else None try: diff --git a/litellm/proxy/_experimental/mcp_server/faults/classify.py b/litellm/proxy/_experimental/mcp_server/faults/classify.py index bb41436b495..7826fc34eac 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/classify.py +++ b/litellm/proxy/_experimental/mcp_server/faults/classify.py @@ -10,6 +10,7 @@ from __future__ import annotations from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.faults.types import ( @@ -25,6 +26,8 @@ from litellm.proxy._experimental.mcp_server.faults.types import ( UpstreamReportedFault, ) +_DECODED_JSON: Final = TypeAdapter(object) + def _safe_text(response: httpx.Response) -> str: try: @@ -35,7 +38,7 @@ def _safe_text(response: httpx.Response) -> str: def _safe_json(response: httpx.Response) -> object: try: - return response.json() + return _DECODED_JSON.validate_python(response.json()) except Exception: return None diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index 84673c8c738..e148e28e826 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -3,8 +3,8 @@ The aggregate listing deliberately keeps serving the healthy subset when one server fails, but a failed server must contribute a classified outcome instead of silently shrinking the list: an empty contribution with no signal makes a broken upstream indistinguishable from a healthy server with no -tools. Outcomes carry only machine fields (category and status code) so nothing from an upstream -body crosses the trust boundary; classification is total, so any exception out of a server fetch +tools. Outcomes carry machine fields and controlled diagnostics so nothing from an upstream body +crosses the trust boundary; classification is total, so any exception out of a server fetch becomes an outcome, never a second failure. """ @@ -21,8 +21,10 @@ from urllib.request import parse_http_list, parse_keqv_list import anyio import httpx import httpx2 +from mcp import MCPError +from mcp.types import CONNECTION_CLOSED from mcp.types import Tool as MCPTool -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, PositiveInt, TypeAdapter from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -65,6 +67,16 @@ class ServerListFault(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) tag: ListFaultCategory status_code: int | None = None + sse_event_limit_bytes: PositiveInt | None = None + + @property + def public_message(self) -> str | None: + if self.sse_event_limit_bytes is None: + return None + return ( + f"Upstream MCP SSE event exceeded the {self.sse_event_limit_bytes} byte per-event limit. " + "Paginate tools/list responses so each event stays below the limit." + ) ServerOutcome: TypeAlias = ServerListOk | ServerListFault @@ -155,6 +167,18 @@ def raise_classified_list_failure( raise MCPServerListError(classify_list_exception(exc), server_name) from exc +def _sse_event_limit(exc: BaseException) -> int | None: + """The SDK replaces decoder exceptions with CONNECTION_CLOSED and a fixed diagnostic.""" + if not isinstance(exc, (httpx2.SSEError, MCPError)): + return None + if isinstance(exc, MCPError) and exc.error.code != CONNECTION_CLOSED: + return None + prefix: Final = "SSE stream failed: " if isinstance(exc, MCPError) else "" + message: Final = exc.error.message if isinstance(exc, MCPError) else str(exc) + matched: Final = re.fullmatch(prefix + r"Server-sent event exceeded the ([1-9][0-9]{0,18}) byte limit\.", message) + return int(matched[1]) if matched is not None else None + + def classify_list_exception(exc: BaseException) -> ServerListFault: """Classify a per-server listing failure into exactly one outcome. Total: an exception this function cannot recognize is the gateway's own fault (``internal``), never a re-raise.""" @@ -179,6 +203,11 @@ def classify_list_exception(exc: BaseException) -> ServerListFault: response: Final = _find_upstream_response(exc) if response is not None: return ServerListFault(tag="upstream_error", status_code=response.status_code) + sse_event_limit: Final = next( + (limit for error in iter_exception_tree(exc) if (limit := _sse_event_limit(error)) is not None), None + ) + if sse_event_limit is not None: + return ServerListFault(tag="upstream_error", sse_event_limit_bytes=sse_event_limit) if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)): return ServerListFault(tag="timeout") if isinstance(exc, (httpx.TransportError, httpx2.TransportError)): @@ -226,7 +255,7 @@ def rate_limit_retry_delay( ) -> float: if retry_after_seconds is not None: return min(retry_after_seconds, max_delay) - return min(max_delay, base_delay * 2**attempt * (1 + jitter)) + return min(max_delay, base_delay * 2.0**attempt * (1 + jitter)) T: Final = TypeVar("T") @@ -295,7 +324,7 @@ def listing_failed_outright(outcomes: Mapping[str, ServerOutcome]) -> bool: def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: """The client-visible form of one outcome, for the tools/list result ``_meta`` and the REST - response: category plus status code only, never upstream prose or URLs.""" + response: category, status code and controlled diagnostics, never upstream prose or URLs.""" match outcome.tag: case "ok": return {"status": "ok", "tool_count": outcome.tool_count} @@ -303,6 +332,15 @@ def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: return { "status": outcome.tag, **({"http_status": outcome.status_code} if outcome.status_code is not None else {}), + **( + { + "reason": "sse_event_too_large", + "limit_bytes": outcome.sse_event_limit_bytes, + "message": outcome.public_message, + } + if outcome.sse_event_limit_bytes is not None + else {} + ), } case _: assert_never(outcome.tag) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7eddf592262..db70c7d77de 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -488,6 +488,10 @@ class _AuthorizationServerMetadataPayload(TypedDict, total=False): token_endpoint_auth_methods_supported: Sequence[str] +def _mapping_value(source: Mapping[str, object], key: str, default: object) -> object: + return source.get(key, default) + + def _blank_to_none(value: str | None) -> str | None: """Collapse an absent, empty, or whitespace-only string to ``None``. @@ -2757,7 +2761,9 @@ class MCPServerManager: server_id=server_id, name=name_for_prefix, protocol_version=TypeAdapter(MCPUpstreamProtocol).validate_python( - server_config.get("protocol_version", mcp_info.get("protocol_version", "auto")) + _mapping_value( + server_config, "protocol_version", _mapping_value(mcp_info, "protocol_version", "auto") + ) ), alias=alias, server_name=server_name, diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 6d4fefcdb6c..cc92b8cce9b 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -11,6 +11,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Final import httpx +from pydantic import TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -41,6 +42,8 @@ from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +_DECODED_JSON: Final = TypeAdapter(object) + class MCPOAuth2TokenCache(InMemoryCache): """ @@ -160,7 +163,7 @@ class MCPOAuth2TokenCache(InMemoryCache): f"failed with status {exc.response.status_code}" ) from exc - body: Final = response.json() + body: Final = _DECODED_JSON.validate_python(response.json()) if not isinstance(body, dict): raise ValueError( diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index c31560a9f63..b3aa040c507 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -18,6 +18,7 @@ from typing import Final, Literal, Protocol, TypeAlias import jwt from fastapi import HTTPException from jwt.types import Options +from pydantic import TypeAdapter from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -38,6 +39,7 @@ _ALLOWED_ID_TOKEN_ALGORITHMS: Final = ( "PS512", ) _JWKS_CACHE_TTL_SECONDS: Final = 3600 +_DECODED_JSON: Final = TypeAdapter(object) _jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS) JwksFetcher: TypeAlias = Callable[ @@ -103,7 +105,7 @@ async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> Sequence[Mappi client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response: Final = await client.get(jwks_url) response.raise_for_status() - document: Final = response.json() + document: Final = _DECODED_JSON.validate_python(response.json()) keys: Final = document.get("keys") if isinstance(document, dict) else None if not isinstance(keys, list): raise TypeError(f"JWKS document at {jwks_url} has no 'keys' array") @@ -116,7 +118,7 @@ async def _discover_jwks_url(issuer: str) -> str: client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response: Final = await client.get(discovery_url) response.raise_for_status() - metadata: Final = response.json() + metadata: Final = _DECODED_JSON.validate_python(response.json()) jwks_uri: Final = metadata.get("jwks_uri") if isinstance(metadata, dict) else None if not isinstance(jwks_uri, str) or not jwks_uri: raise ValueError(f"OIDC discovery at {discovery_url} returned no jwks_uri") diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py index 2501f751d8e..d84a65fee02 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py @@ -12,6 +12,7 @@ from __future__ import annotations from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( @@ -34,6 +35,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger _GATEWAY_FAULT_OAUTH_ERRORS: Final = frozenset( {"invalid_client", "unauthorized_client", "unsupported_grant_type", "invalid_target", "invalid_scope"} ) +_DECODED_JSON: Final = TypeAdapter(object) def _oauth_error_fields(response: httpx.Response) -> tuple[str | None, str | None]: @@ -46,7 +48,7 @@ def _oauth_error_fields(response: httpx.Response) -> tuple[str | None, str | Non it can carry IdP internals and must never reach the caller. """ try: - body: Final[object] = response.json() + body: Final = _DECODED_JSON.validate_python(response.json()) except Exception: # noqa: BLE001 return None, None if not isinstance(body, dict): diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index ee063b81ae4..62387b7cc28 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -957,7 +957,7 @@ if MCP_AVAILABLE: status_code=list_fault_http_status(fault), detail={ "error": fault.tag, - "message": f"Failed to list tools from server {get_server_prefix(server)}", + "message": fault.public_message or f"Failed to list tools from server {get_server_prefix(server)}", }, ) from e except Exception as e: diff --git a/litellm/proxy/_experimental/mcp_server/result_conversion.py b/litellm/proxy/_experimental/mcp_server/result_conversion.py index 61435f94c94..fdb9a01fea6 100644 --- a/litellm/proxy/_experimental/mcp_server/result_conversion.py +++ b/litellm/proxy/_experimental/mcp_server/result_conversion.py @@ -13,7 +13,7 @@ from __future__ import annotations import json import math from collections.abc import Sequence -from typing import Final, TypeAlias, TypeVar +from typing import Final, Protocol, TypeAlias, TypeVar from mcp.types import CacheableResult, CallToolResult, ContentBlock, InputRequiredResult, TextContent, Tool from typing_extensions import ReadOnly, TypedDict, assert_never @@ -105,8 +105,17 @@ def complete_call_tool_result(outcome: ToolOutcome, compat: WireCompat) -> CallT return converted +class _StructuredContentCarrier(Protocol): + @property + def structured_content(self) -> object: ... + + +def _structured_content(result: _StructuredContentCarrier) -> object: + return result.structured_content + + def _downgrade_structured_content(result: CallToolResult) -> CallToolResult: - structured: Final = result.structured_content + structured: Final = _structured_content(result) if structured is None or isinstance(structured, dict): return result fallback: Final = TextContent(type="text", text=json.dumps(structured)) diff --git a/litellm/proxy/_experimental/mcp_server/scope_challenge.py b/litellm/proxy/_experimental/mcp_server/scope_challenge.py index c326418bee7..121234e305d 100644 --- a/litellm/proxy/_experimental/mcp_server/scope_challenge.py +++ b/litellm/proxy/_experimental/mcp_server/scope_challenge.py @@ -33,9 +33,8 @@ class OAuthScopeResponse: self._request_path = request_path self._ready = asyncio.Event() self._lock = asyncio.Lock() - self._start: Message | None = None self._challenge: MCPUpstreamAuthError | None = None - self._refused = False + self._closed = False @property def pending(self) -> bool: @@ -50,31 +49,33 @@ class OAuthScopeResponse: self._ready.set() async def __call__(self, message: Message) -> None: - if message["type"] == "http.response.start": - self._start = message - if message["status"] != 200: - self.allow() - return - await self._ready.wait() async with self._lock: - if self._refused: + if self._closed: return - if self._start is not None: - start: Final = self._start - self._start = None - if self._challenge is not None: - error: Final = self._challenge.to_http_exception(self._base_url, self._request_path) - response: Final = JSONResponse( - status_code=error.status_code, content={"detail": error.detail}, headers=error.headers - ) - self._refused = True - await self._send( - {"type": "http.response.start", "status": response.status_code, "headers": response.raw_headers} - ) - await self._send({"type": "http.response.body", "body": bytes(response.body)}) - return - await self._send(start) - await self._send(message) + try: + if message["type"] == "http.response.start": + if message["status"] != 200: + self.allow() + await self._ready.wait() + if self._challenge is not None: + error: Final = self._challenge.to_http_exception(self._base_url, self._request_path) + response: Final = JSONResponse( + status_code=error.status_code, content={"detail": error.detail}, headers=error.headers + ) + self._closed = True + await self._send( + { + "type": "http.response.start", + "status": response.status_code, + "headers": response.raw_headers, + } + ) + await self._send({"type": "http.response.body", "body": bytes(response.body)}) + return + await self._send(message) + except BaseException: + self._closed = True + raise def scope_response_from_request(request: object) -> OAuthScopeResponse | None: diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 2f88d8c361f..a7748b9ca88 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -6,7 +6,7 @@ Filters MCP tools semantically for /chat/completions and /responses endpoints. import asyncio from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, TypeVar from litellm._logging import verbose_logger from litellm.exceptions import ContextWindowExceededError @@ -20,6 +20,8 @@ if TYPE_CHECKING: from litellm.router import Router +_Tool: Final = TypeVar("_Tool") + class SemanticToolFilterContextWindowError(Exception): """Raised when the embedding model exceeds its context window, so semantic filtering cannot run.""" @@ -265,9 +267,9 @@ class SemanticMCPToolFilter: async def filter_tools( self, query: str, - available_tools: list[Any], + available_tools: list[_Tool], top_k: int | None = None, - ) -> list[Any]: + ) -> list[_Tool]: """ Filter tools semantically based on query. @@ -390,7 +392,7 @@ class SemanticMCPToolFilter: separator: Final = client_name[-len(canonical) - 1] return separator in ("_", "-") - def _get_tools_by_names(self, tool_names: Sequence[str], available_tools: Sequence[object]) -> list[object]: + def _get_tools_by_names(self, tool_names: Sequence[str], available_tools: Sequence[_Tool]) -> list[_Tool]: """ Get tools from available_tools by their names, preserving the semantic router's ordering. @@ -406,13 +408,13 @@ class SemanticMCPToolFilter: # Exact matches win over suffix matches when both are present, and # each incoming tool is returned at most once even if two canonical # names happen to be tail-compatible with the same incoming name. - available_by_name: Final[dict[str, object]] = {} + available_by_name: Final[dict[str, _Tool]] = {} for tool in available_tools: client_name, _ = self.extract_tool_info(tool) if client_name and client_name not in available_by_name: available_by_name[client_name] = tool - matched: Final[list[object]] = [] + matched: Final[list[_Tool]] = [] used_ids: Final[set[int]] = set() for canonical in tool_names: tool = available_by_name.get(canonical) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b3426089375..7207c4150ad 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -14,7 +14,7 @@ import time import types from collections import Counter from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence, Set -from typing import TYPE_CHECKING, Final, NoReturn, Protocol +from typing import TYPE_CHECKING, Final, Literal, NoReturn, Protocol import httpx from fastapi import FastAPI, HTTPException @@ -150,6 +150,16 @@ def _scope_headers(scope: Scope) -> Sequence[tuple[bytes, bytes]]: return asgi_scope_fields(scope).get("headers", []) +class _RawHeaderPairsSource(Protocol): + def get( + self, key: Literal["headers"], default: Sequence[tuple[bytes | str, bytes | str]], / + ) -> Sequence[tuple[bytes | str, bytes | str]]: ... + + +def _raw_header_pairs(source: _RawHeaderPairsSource) -> Sequence[tuple[bytes | str, bytes | str]]: + return source.get("headers", []) + + def unsupported_protocol_version(scope: Scope) -> str | None: """Admit configured HTTP revisions while keeping SSE on the legacy protocol.""" from litellm.proxy._experimental.mcp_server.capabilities import configured_versions @@ -1319,7 +1329,7 @@ if MCP_AVAILABLE: Extract mcp-session-id from ASGI scope headers. Returns None if not present. """ - scope_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) + scope_headers: Final = _raw_header_pairs(scope) for header_name, header_value in scope_headers: name = header_name if isinstance(header_name, bytes) else header_name.encode() if name.lower() == b"mcp-session-id": @@ -1669,7 +1679,7 @@ if MCP_AVAILABLE: Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header: Final = b"mcp-session-id" - _headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) + _headers: Final = _raw_header_pairs(scope) def _normalize_header_name(header_name: object) -> bytes | None: if isinstance(header_name, bytes): @@ -1797,7 +1807,7 @@ if MCP_AVAILABLE: return acting.model_copy(update={"mcp_toolset_id": toolset_id}) tool_permissions = await operations.global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=[toolset_id] + toolset_ids=[toolset_id], requires_fresh_policy=acting.requires_fresh_policy ) server_ids: Final = list(tool_permissions.keys()) existing_op: Final = acting.object_permission @@ -2869,7 +2879,7 @@ if MCP_AVAILABLE: mcp_servers, _sse_client_ip, scoped_server_endpoint=scoped_server_endpoint, - is_initialize=scope.get("method") == "GET", + is_initialize=asgi_scope_fields(scope).get("method", "") == "GET", ): async with ( sse.connect_sse(transport_scope, receive, send) as (read_stream, write_stream), @@ -3090,7 +3100,7 @@ if MCP_AVAILABLE: ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": - response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", []) + response_headers: Final = _raw_header_pairs(message) for key, value in response_headers: header_name = key if isinstance(key, bytes) else str(key).encode() if header_name.lower() == b"mcp-session-id": diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index f3a78d206d7..35906c793b4 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -131,6 +131,10 @@ async def get_mcp_toolset_by_name( return _toolset_from_row(row) +def _set_fields(data: UpdateMCPToolsetRequest) -> Mapping[str, object]: + return data.model_dump(exclude_unset=True) + + async def update_mcp_toolset( prisma_client: PrismaClient, data: UpdateMCPToolsetRequest, @@ -143,7 +147,7 @@ async def update_mcp_toolset( data_dict: Final = dict( ( (field, json.dumps(value) if field == "tools" else value) - for field, value in data.model_dump(exclude_unset=True).items() + for field, value in _set_fields(data).items() if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None) ), updated_by=touched_by, diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 1684b9c5c30..ee8f3ea3b30 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -13,6 +13,8 @@ from collections.abc import Set as AbstractSet from typing import Any, Final, Literal, Protocol, overload from urllib.parse import quote +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -21,6 +23,7 @@ if typing.TYPE_CHECKING: MCP_SERVERS_TARGET: Final = "mcp_servers" MCP_OAUTH_TOKENS_TARGET: Final = "mcp_oauth_tokens" +_DECODED_JSON: Final = TypeAdapter(object) class ASGIScopeFields(Protocol): @@ -243,7 +246,7 @@ def _parse_mcp_info_dict(mcp_info: object) -> Mapping[str, object] | None: return mcp_info if isinstance(mcp_info, str): try: - parsed: Final[object] = json.loads(mcp_info) + parsed: Final = _DECODED_JSON.validate_python(json.loads(mcp_info)) except (ValueError, TypeError): return None return parsed if isinstance(parsed, dict) else None diff --git a/litellm/proxy/_experimental/out/assets/aria-agent.png b/litellm/proxy/_experimental/out/assets/aria-agent.png new file mode 100644 index 00000000000..870773933e3 Binary files /dev/null and b/litellm/proxy/_experimental/out/assets/aria-agent.png differ diff --git a/litellm/proxy/_experimental/out/assets/logos/claude.svg b/litellm/proxy/_experimental/out/assets/logos/claude.svg new file mode 100644 index 00000000000..ede30652b25 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/claude.svg @@ -0,0 +1,27 @@ + +Claude diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 0e2eed73364..4e54d6ac733 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -207,6 +207,11 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.management_endpoints.config_override_endpoints", path_prefixes=("/config_overrides",), ), + LazyFeature( + name="codex_passthrough", + module_path="litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints", + path_prefixes=("/backend-api/codex/",), + ), LazyFeature( name="llm_passthrough", module_path="litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints", @@ -299,11 +304,6 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.spend_tracking.vantage_endpoints", path_prefixes=("/vantage",), ), - LazyFeature( - name="usage_ai", - module_path="litellm.proxy.management_endpoints.usage_endpoints", - path_prefixes=("/usage/ai",), - ), LazyFeature( name="prompts", module_path="litellm.proxy.prompts.prompt_endpoints", @@ -361,7 +361,7 @@ class LazyFeatureMiddleware: async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: # Short-circuit once every feature has loaded. if scope["type"] in ("http", "websocket") and len(self._loaded) < len(self._features): - path: str = scope.get("path", "") + path: str = str(_scope_value(scope, "path")) # Strip the request's root_path so prefix matching works under a # server root path. Without this, requests like /api/v1/policies/... # never match the registered prefixes (/policies/...) and lazy @@ -589,6 +589,7 @@ def _make_warmup_router(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY from fastapi import APIRouter, Depends, HTTPException from fastapi.openapi.utils import get_openapi + from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router: Final = APIRouter() @@ -608,17 +609,22 @@ def _make_warmup_router(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY await _force_load(app, feat, features) feat_routes: Final = [r for r in app.routes if feat.matches(getattr(r, "path", ""))] - full: Final = get_openapi(title=app.title, version=app.version, routes=feat_routes) + full: Final[Mapping[str, object]] = MappingProxyType( + get_openapi(title=app.title, version=app.version, routes=feat_routes) + ) + paths: Final = as_str_mapping(full.get("paths", {})) or {} + components: Final = as_str_mapping(full.get("components", {})) or {} # Force all operations under one tag so they group under a single Swagger # section — many lazy modules tag routes inconsistently. - for path_ops in full.get("paths", {}).values(): - for op in path_ops.values(): - if isinstance(op, dict): - op["tags"] = [feat.name] + for path_ops in paths.values(): + for op in (as_str_mapping(path_ops) or {}).values(): + operation = as_str_mapping(op) + if isinstance(operation, dict): + operation["tags"] = [feat.name] return { "stub_path": feat.path_prefixes[0], "paths": full.get("paths", {}), - "components": {"schemas": full.get("components", {}).get("schemas", {})}, + "components": {"schemas": components.get("schemas", {})}, } return router diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 545b317055a..166072e9613 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -8673,6 +8673,184 @@ } } }, + "codex_passthrough": { + "components": { + "schemas": { + "HTTPValidationError": { + "properties": { + "detail": { + "items": { + "$ref": "#/components/schemas/ValidationError" + }, + "title": "Detail", + "type": "array" + } + }, + "title": "HTTPValidationError", + "type": "object" + }, + "ValidationError": { + "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, + "loc": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "title": "Location", + "type": "array" + }, + "msg": { + "title": "Message", + "type": "string" + }, + "type": { + "title": "Error Type", + "type": "string" + } + }, + "required": [ + "loc", + "msg", + "type" + ], + "title": "ValidationError", + "type": "object" + } + } + }, + "paths": { + "/backend-api/codex/{endpoint}": { + "get": { + "description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication.", + "operationId": "codex_passthrough_route_backend_api_codex__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "enum": [ + "models", + "responses" + ], + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Codex Passthrough Route", + "tags": [ + "codex_passthrough" + ] + }, + "post": { + "description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication.", + "operationId": "codex_passthrough_route_backend_api_codex__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "enum": [ + "models", + "responses" + ], + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "text/event-stream": { + "schema": { + "type": "string" + } + } + }, + "description": "Unmodified Codex response or server-sent events" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Codex Passthrough Route", + "tags": [ + "codex_passthrough" + ] + } + } + } + }, "compliance": { "components": { "schemas": { @@ -11407,7 +11585,7 @@ }, "/systemone": { "post": { - "description": "Judge a `state` against 1 to 128 named `questions` and get one calibrated answer per question name.\n\nQuestion types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with\n`confidence` and `probabilities`), `score` (an index into the `criteria` levels, with `confidence`,\n`legend` and `probabilities`) and `box` (the image region the `instructions` describe, with `confidence`).\nTypeSafe models also take images in `assets`, each one referenced by its key in `state`. Requests that break\nthese rules are rejected with a 400 before any provider is called. `model` is any decision model in the proxy model list. OpenAI decision models accept this format\ntoo, LiteLLM translates it and keeps the answers keyed by question name. Streaming is not supported.\n\n[Docs](https://docs.litellm.ai/docs/decisions)", + "description": "Judge a `state` against one or more named `questions` and get one calibrated answer per question name.\n\nQuestion types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with\n`confidence` and `probabilities`), `score` (an index into the `criteria` levels, with `confidence`,\n`legend` and `probabilities`) and `box` (the image region the `instructions` describe, with `confidence`).\nTypeSafe models also take images in `assets`, each one referenced by its key in `state`. Requests that break\nthese rules are rejected with a 400 before any provider is called. `model` is any decision model in the proxy model list. OpenAI decision models accept this format\ntoo, LiteLLM translates it and keeps the answers keyed by question name. Streaming is not supported.\n\n[Docs](https://docs.litellm.ai/docs/decisions)", "operationId": "systemone_systemone_post", "requestBody": { "content": { @@ -11813,7 +11991,6 @@ ] }, "description": "Named noul, choice and score questions. Each key becomes a key in the answers", - "maxProperties": 128, "minProperties": 1, "propertyNames": { "minLength": 1 @@ -12814,7 +12991,7 @@ }, "/v1/systemone": { "post": { - "description": "Judge a `state` against 1 to 128 named `questions` and get one calibrated answer per question name.\n\nQuestion types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with\n`confidence` and `probabilities`), `score` (an index into the `criteria` levels, with `confidence`,\n`legend` and `probabilities`) and `box` (the image region the `instructions` describe, with `confidence`).\nTypeSafe models also take images in `assets`, each one referenced by its key in `state`. Requests that break\nthese rules are rejected with a 400 before any provider is called. `model` is any decision model in the proxy model list. OpenAI decision models accept this format\ntoo, LiteLLM translates it and keeps the answers keyed by question name. Streaming is not supported.\n\n[Docs](https://docs.litellm.ai/docs/decisions)", + "description": "Judge a `state` against one or more named `questions` and get one calibrated answer per question name.\n\nQuestion types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with\n`confidence` and `probabilities`), `score` (an index into the `criteria` levels, with `confidence`,\n`legend` and `probabilities`) and `box` (the image region the `instructions` describe, with `confidence`).\nTypeSafe models also take images in `assets`, each one referenced by its key in `state`. Requests that break\nthese rules are rejected with a 400 before any provider is called. `model` is any decision model in the proxy model list. OpenAI decision models accept this format\ntoo, LiteLLM translates it and keeps the answers keyed by question name. Streaming is not supported.\n\n[Docs](https://docs.litellm.ai/docs/decisions)", "operationId": "systemone_v1_systemone_post", "requestBody": { "content": { @@ -13220,7 +13397,6 @@ ] }, "description": "Named noul, choice and score questions. Each key becomes a key in the answers", - "maxProperties": 128, "minProperties": 1, "propertyNames": { "minLength": 1 @@ -22860,6 +23036,25 @@ "title": "Body_image_edit_api_openai_deployments__model__images_edits_post", "type": "object" }, + "CacheWriteTokensDetails": { + "additionalProperties": true, + "properties": { + "ttl_12h_tokens": { + "title": "Ttl 12H Tokens", + "type": "integer" + }, + "ttl_30m_tokens": { + "title": "Ttl 30M Tokens", + "type": "integer" + } + }, + "required": [ + "ttl_30m_tokens", + "ttl_12h_tokens" + ], + "title": "CacheWriteTokensDetails", + "type": "object" + }, "CachedTokensDetails": { "properties": { "audio_tokens": { @@ -24354,6 +24549,27 @@ ], "title": "Audio Tokens" }, + "cache_write_tokens": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Cache Write Tokens" + }, + "cache_write_tokens_details": { + "anyOf": [ + { + "$ref": "#/components/schemas/CacheWriteTokensDetails" + }, + { + "type": "null" + } + ] + }, "cached_tokens": { "default": 0, "title": "Cached Tokens", @@ -28240,6 +28456,9 @@ }, { "$ref": "#/components/schemas/OutputToolSearchCall" + }, + { + "$ref": "#/components/schemas/ResponseCompactionItem" } ] }, @@ -57897,163 +58116,6 @@ } } }, - "usage_ai": { - "components": { - "schemas": { - "ChatMessage": { - "properties": { - "content": { - "title": "Content", - "type": "string" - }, - "role": { - "enum": [ - "user", - "assistant" - ], - "title": "Role", - "type": "string" - } - }, - "required": [ - "role", - "content" - ], - "title": "ChatMessage", - "type": "object" - }, - "HTTPValidationError": { - "properties": { - "detail": { - "items": { - "$ref": "#/components/schemas/ValidationError" - }, - "title": "Detail", - "type": "array" - } - }, - "title": "HTTPValidationError", - "type": "object" - }, - "UsageAIChatRequest": { - "properties": { - "messages": { - "description": "Chat messages (user/assistant history)", - "items": { - "$ref": "#/components/schemas/ChatMessage" - }, - "title": "Messages", - "type": "array" - }, - "model": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "description": "Model to use for AI chat", - "title": "Model" - } - }, - "required": [ - "messages" - ], - "title": "UsageAIChatRequest", - "type": "object" - }, - "ValidationError": { - "properties": { - "ctx": { - "title": "Context", - "type": "object" - }, - "input": { - "title": "Input" - }, - "loc": { - "items": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "integer" - } - ] - }, - "title": "Location", - "type": "array" - }, - "msg": { - "title": "Message", - "type": "string" - }, - "type": { - "title": "Error Type", - "type": "string" - } - }, - "required": [ - "loc", - "msg", - "type" - ], - "title": "ValidationError", - "type": "object" - } - } - }, - "paths": { - "/usage/ai/chat": { - "post": { - "description": "AI chat about usage data. Streams SSE events with the AI response.\nThe AI agent has access to tools that query aggregated daily activity data.", - "operationId": "usage_ai_chat_usage_ai_chat_post", - "requestBody": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/UsageAIChatRequest" - } - } - }, - "required": true - }, - "responses": { - "200": { - "content": { - "application/json": { - "schema": {} - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Usage Ai Chat", - "tags": [ - "usage_ai" - ] - } - } - } - }, "vantage": { "components": { "schemas": { diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index 00a05307a89..c5a63b0f990 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -16,6 +16,7 @@ import sys from collections.abc import Callable, Mapping from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import TYPE_CHECKING, Final from typing_extensions import ReadOnly, TypedDict @@ -61,7 +62,7 @@ def load_snapshot() -> dict[str, dict] | None: return None -def _normalize_operation_ids(paths: dict[str, dict]) -> None: +def _normalize_operation_ids(paths: Mapping[str, object]) -> None: """Make FastAPI-generated operation IDs stable for multi-method routes. FastAPI derives the default operation ID suffix from the first item in the @@ -69,26 +70,30 @@ def _normalize_operation_ids(paths: dict[str, dict]) -> None: set iteration order can vary between processes, which makes the snapshot drift even when no routes changed. """ + from litellm.litellm_core_utils.core_helpers import as_str_mapping + for path_ops in paths.values(): - if not isinstance(path_ops, dict): + operations = as_str_mapping(path_ops) + if not isinstance(operations, dict): continue - methods = {method for method in path_ops if method in HTTP_METHOD_SUFFIXES} + methods = {method for method in operations if method in HTTP_METHOD_SUFFIXES} if not methods: continue - for method, operation in path_ops.items(): - if method not in HTTP_METHOD_SUFFIXES or not isinstance(operation, dict): + for method, operation in operations.items(): + entry = as_str_mapping(operation) + if method not in HTTP_METHOD_SUFFIXES or not isinstance(entry, dict): continue - operation_id = operation.get("operationId") + operation_id = entry.get("operationId") if not isinstance(operation_id, str): continue for suffix in methods: suffix_token = f"_{suffix}" if operation_id.endswith(suffix_token): - operation["operationId"] = operation_id[: -len(suffix_token)] + f"_{method}" + entry["operationId"] = operation_id[: -len(suffix_token)] + f"_{method}" break @@ -117,6 +122,7 @@ def _register_feature(app: "FastAPI", feat: "LazyFeature") -> str | None: def _feature_fragment(app: "FastAPI", feat: "LazyFeature", used_operation_ids: set[str]) -> SnapshotFragment | None: from fastapi.openapi.utils import get_openapi + from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.proxy_server import ensure_unique_openapi_operation_ids feat_routes: Final = [r for r in app.routes if feat.matches(getattr(r, "path", ""))] @@ -124,22 +130,29 @@ def _feature_fragment(app: "FastAPI", feat: "LazyFeature", used_operation_ids: s return None _stabilize_multi_method_route_ids(feat_routes) full: Final = get_openapi(title=app.title, version=app.version, routes=feat_routes) - paths: Final = full.get("paths", {}) + openapi: Final[Mapping[str, object]] = MappingProxyType(full) + paths: Final = as_str_mapping(openapi.get("paths", {})) or {} _normalize_operation_ids(paths) for path_ops in paths.values(): - for method, op in path_ops.items(): - if isinstance(op, dict): - operation_id = op.get("operationId") + for method, op in (as_str_mapping(path_ops) or {}).items(): + operation = as_str_mapping(op) + if isinstance(operation, dict): + operation_id = operation.get("operationId") if isinstance(operation_id, str): for suffix in HTTP_METHOD_SUFFIXES: if operation_id.endswith(f"_{suffix}"): - op["operationId"] = operation_id[: -len(suffix)] + method + operation["operationId"] = operation_id[: -len(suffix)] + method break - op["tags"] = [feat.name] - unique: Final = ensure_unique_openapi_operation_ids(full, used_operation_ids) + operation["tags"] = [feat.name] + unique: Final[Mapping[str, object]] = MappingProxyType( + ensure_unique_openapi_operation_ids(full, used_operation_ids) + ) + components: Final = as_str_mapping(unique.get("components", {})) or {} return { - "paths": paths, - "components": {"schemas": unique.get("components", {}).get("schemas", {})}, + "paths": { + path: operations for path, entry in paths.items() if (operations := as_str_mapping(entry)) is not None + }, + "components": {"schemas": as_str_mapping(components.get("schemas", {})) or {}}, } diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 54acf59ef80..6280e2caf87 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -219,6 +219,8 @@ from litellm.types.proxy.auth.user_roles import ( # noqa: E402 # re-export aft KeyManagementRoutes as KeyManagementRoutes, # noqa: PLC0414 # public re-export ) +MODEL_LIST_ROUTES: Final = ("/models", "/v1/models", "/claude_code_gateway/v1/models") + class LiteLLMRoutes(enum.Enum): openai_route_names = [ @@ -307,8 +309,7 @@ class LiteLLMRoutes(enum.Enum): "/threads/{thread_id}/runs", "/v1/threads/{thread_id}/runs", # models - "/models", - "/v1/models", + *MODEL_LIST_ROUTES, # token counter "/utils/token_counter", "/utils/model_info", @@ -390,6 +391,7 @@ class LiteLLMRoutes(enum.Enum): ] mapped_pass_through_routes = [ + "/backend-api/codex", "/bedrock", "/comprehendmedical", "/azure_speech", @@ -577,8 +579,7 @@ class LiteLLMRoutes(enum.Enum): "/health/services", "/key/list", "/user/filter/ui", - "/models", - "/v1/models", + *MODEL_LIST_ROUTES, "/sso/get/ui_settings", "/get/user_banner", "/get/latest_release_info", @@ -760,8 +761,7 @@ class LiteLLMRoutes(enum.Enum): "/model/info", "/v2/model/info", "/v2/key/info", - "/models", - "/v1/models", + *MODEL_LIST_ROUTES, "/global/spend", "/global/spend/logs", "/global/spend/keys", @@ -1767,8 +1767,10 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase): max_parallel_requests: int | None = Field( default=None, description="Max concurrent requests allowed for this budget id." ) - tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.") - rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.") + tpm_limit: int | None = Field(default=None, ge=0, description="Max tokens per minute, allowed for this budget id.") + rpm_limit: int | None = Field( + default=None, ge=0, description="Max requests per minute, allowed for this budget id." + ) tpd_limit: int | None = Field( default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id." ) @@ -1879,6 +1881,8 @@ RouterSettingsDict = Annotated[ class NewTeamRequest(TeamBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) router_settings: RouterSettingsDict | None = None model_aliases: dict | None = None model_max_budget: GenericBudgetConfigType | None = Field( @@ -1958,8 +1962,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): team_alias: str | None = None organization_id: str | None = None metadata: dict | None = None - tpm_limit: int | None = None - rpm_limit: int | None = None + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None @@ -2717,6 +2721,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "hidden model can still be called." ), ) + return_wildcard_routes: bool | None = Field( + None, + description=( + "When true, `/v1/models` lists wildcard routes like `openai/*` unless the request passes " + "`return_wildcard_routes=false`." + ), + ) alerting: list | None = Field( None, description="List of alerting integrations - e.g. `alerting: ['slack', 'webhook', 'email']`. 'slack' posts Slack-format messages to any Slack-compatible webhook (Slack, Rocket.Chat, Mattermost); 'webhook' posts structured JSON budget alerts to WEBHOOK_URL", @@ -5255,6 +5266,7 @@ class ToolDiscoveryQueueItem(TypedDict, total=False): team_id: str | None # team that triggered discovery key_alias: str | None # human-readable key alias user_agent: str | None # HTTP User-Agent of the caller + call_count: ReadOnly[int] # invocations folded into this row since the previous flush from litellm.models.managed_files import ( # noqa: E402 diff --git a/litellm/proxy/a2a/discovery.py b/litellm/proxy/a2a/discovery.py index ff58c9c85ec..03d1be7d185 100644 --- a/litellm/proxy/a2a/discovery.py +++ b/litellm/proxy/a2a/discovery.py @@ -16,10 +16,13 @@ fetcher dispatches by ``discovery_mode``: from collections.abc import Mapping from enum import Enum -from typing import Any, Final +from typing import Final from urllib.parse import urlencode +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -46,6 +49,8 @@ AGENT_CARD_WELL_KNOWN_PATHS: Final[tuple[str, ...]] = ( DEFAULT_DISCOVERY_TIMEOUT_SECONDS: Final = 10.0 +_DECODED_JSON: Final = TypeAdapter(object) + class AgentCardDiscoveryError(Exception): """Raised when none of the well-known paths returned a usable agent card.""" @@ -87,7 +92,7 @@ async def fetch_well_known_card( params: Mapping[str, object] | None = None, timeout: float = DEFAULT_DISCOVERY_TIMEOUT_SECONDS, headers: dict[str, str] | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Fetch an agent card from ``base_url`` using the strategy chosen by ``discovery_mode``. Returns the parsed JSON from the first path that @@ -137,13 +142,14 @@ async def fetch_well_known_card( continue try: - card = response.json() + decoded = _DECODED_JSON.validate_python(response.json()) except Exception as exc: last_error = f"{url}: invalid JSON ({exc})" continue + card = as_str_mapping(decoded) if not isinstance(card, dict): - last_error = f"{url}: expected JSON object, got {type(card).__name__}" + last_error = f"{url}: expected JSON object, got {type(decoded).__name__}" continue verbose_proxy_logger.debug("A2A discovery succeeded at %s", url) diff --git a/litellm/proxy/a2a/endpoints.py b/litellm/proxy/a2a/endpoints.py index 56b512df934..730d40b46d0 100644 --- a/litellm/proxy/a2a/endpoints.py +++ b/litellm/proxy/a2a/endpoints.py @@ -9,7 +9,7 @@ admin pick which ones to expose through the proxy. The actual merge into a LiteLLM-fronted card happens when the agent is saved via ``POST /v1/agents``. """ -from typing import Any, Final +from typing import Final from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse @@ -50,7 +50,7 @@ class DiscoverAgentRequest(LiteLLMBaseModel): "query parameter." ), ) - params: dict[str, Any] | None = Field( + params: dict[str, object] | None = Field( default=None, description=( "Mode-specific parameters. ``langgraph_platform`` requires " @@ -61,7 +61,7 @@ class DiscoverAgentRequest(LiteLLMBaseModel): class DiscoverAgentResponse(LiteLLMBaseModel): url: str - agent_card: dict[str, Any] + agent_card: dict[str, object] @router.post( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index c11d296837e..4691c7b0bc5 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -23,6 +23,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header from litellm.proxy._types import UserAPIKeyAuth @@ -281,7 +282,7 @@ async def _a2a_sse_event_source( request_id: str | int | None = None, extra_headers: Mapping[str, str] | None = None, served_version: A2AVersion = "0.3", -) -> AsyncGenerator[Mapping[str, object], None]: +) -> AsyncGenerator[object, None]: """Stream an upstream A2A SSE response as parsed JSON-RPC event dicts. Upstream HTTP/JSON-RPC errors are surfaced as a single JSON-RPC error event @@ -327,12 +328,14 @@ async def _a2a_sse_event_source( if not payload: continue try: - event = json.loads(payload) + event = _DECODED_JSON.validate_python(json.loads(payload)) except Exception: continue - if isinstance(event, dict): - event = normalize_stream_event(event, served_version, request_id=request_id) - yield event + event_fields = as_str_mapping(event) + if isinstance(event_fields, dict): + yield normalize_stream_event(event_fields, served_version, request_id=request_id) + else: + yield event finally: await resp.aclose() @@ -741,7 +744,7 @@ async def invoke_agent_a2a( # Get backend URL and agent name agent_card_params: Final = agent.agent_card_params or {} agent_url: Final = agent_card_params.get("url") - agent_name: Final = agent_card_params.get("name", agent_id) + agent_name: Final = (as_str_mapping(agent_card_params) or {}).get("name", agent_id) # Get litellm_params (may include custom_llm_provider for completion bridge) litellm_params: dict[str, object] = agent.litellm_params or {} diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 9a6f674db2c..fb727367caf 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -18,6 +18,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.handler import ( bridge_model_name, ) from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.databricks_oauth import ( resolve_databricks_backend_auth, @@ -173,13 +174,14 @@ async def route_a2a_agent_request( return _call_route(route_type, _bridge_agent_request(data, agent, bridge_auth.litellm_params)) # Get API base URL from agent config - if not agent.agent_card_params or "url" not in agent.agent_card_params: + card_params: Final = as_str_mapping(agent.agent_card_params) + if not card_params or "url" not in card_params: verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) # Inject API base and route to litellm - data["api_base"] = agent.agent_card_params["url"] + data["api_base"] = card_params["url"] verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, data["api_base"]) return _call_route(route_type, data) diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py index 7425c240d10..bebba009d82 100644 --- a/litellm/proxy/agent_endpoints/databricks_oauth.py +++ b/litellm/proxy/agent_endpoints/databricks_oauth.py @@ -48,6 +48,7 @@ _DEFAULT_SCOPE: Final = "all-apis" _TOKEN_EXPIRY_BUFFER_SECONDS: Final = 60 _DEFAULT_TTL_SECONDS: Final = 3600 _OAUTH_BLOCK: Final = TypeAdapter(Mapping[str, object]) +_TOKEN_RESPONSE_BODY: Final = TypeAdapter(object) def _resolve_secret(value: object) -> str | None: @@ -275,7 +276,7 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): except httpx.HTTPError as exc: raise ValueError(f"Databricks App OAuth token request failed: {exc}") from exc - body: Final[object] = response.json() + body: Final = _TOKEN_RESPONSE_BODY.validate_python(response.json()) if not isinstance(body, dict): raise ValueError( f"Databricks App OAuth token response returned non-object JSON (got {type(body).__name__})" diff --git a/litellm/proxy/agent_endpoints/managed_identity.py b/litellm/proxy/agent_endpoints/managed_identity.py index abab21901ee..fab7ebbc61d 100644 --- a/litellm/proxy/agent_endpoints/managed_identity.py +++ b/litellm/proxy/agent_endpoints/managed_identity.py @@ -119,6 +119,10 @@ def managed_write_fields( return AgentIdentityFailure(message=f"Invalid agent identity configuration: {exc}") +def _identity_fields(identity: EntraIdentityConfig) -> Mapping[str, object]: + return identity.model_dump() + + def _identity_write(identity: EntraIdentityConfig | None, existing: AgentResponse | None) -> ManagedWriteFields: if identity is None: unbind: Final[ManagedWriteFields] = { @@ -134,7 +138,7 @@ def _identity_write(identity: EntraIdentityConfig | None, existing: AgentRespons existing and existing.identity and existing.identity.active - and all(getattr(existing.identity, name) == value for name, value in identity.model_dump().items()) + and all(getattr(existing.identity, name) == value for name, value in _identity_fields(identity).items()) ): unchanged: Final[ManagedWriteFields] = {} return unchanged diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 8e348423837..b0088d93112 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -21,6 +21,7 @@ import json import re from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import Annotated, Final, Protocol, TypedDict, get_args from urllib.parse import urlsplit @@ -299,7 +300,7 @@ def _validate_installation_preference(spec: PluginSpec) -> None: def _build_plugin_manifest(name: str, spec: PluginSpec) -> Mapping[str, object]: """Build the stored manifest dict shared by plugin create and update.""" - dumped: Final[Mapping[str, object]] = spec.model_dump(exclude_none=True) + dumped: Final[Mapping[str, object]] = MappingProxyType(spec.model_dump(exclude_none=True)) return {"name": name, **{key: value for key, value in dumped.items() if value and key != "name"}} diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 4c4155379cd..ba3307bc1ad 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -58,6 +58,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg ) from litellm.proxy.management_endpoints.sso_helper_utils import CLI_SSO_SESSIONS_TARGET from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail +from litellm.proxy.proxy_server import model_list from litellm.types.llms.base import LiteLLMBaseModel GATEWAY_PREFIX: Final = "/claude_code_gateway" @@ -158,6 +159,13 @@ router: Final = APIRouter( _GATEWAY_ENABLED: Final = (Depends(ensure_gateway_enabled),) _AUTHENTICATED: Final = (Depends(user_api_key_auth),) +router.add_api_route( + "/v1/models", + model_list, + methods=["GET"], + dependencies=_GATEWAY_ENABLED, + include_in_schema=False, +) router.add_api_route( "/v1/messages", anthropic_response, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8887d86570a..448643ef59b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,6 +15,7 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from functools import partial from itertools import chain from types import MappingProxyType @@ -47,6 +48,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.models.project import LiteLLM_ProjectTable from litellm.proxy._types import ( + MODEL_LIST_ROUTES, RBAC_ROLES, UI_TEAM_ID, CallInfo, @@ -393,7 +395,7 @@ def _member_allowed_models(membership: _TeamMembershipModelScope) -> Sequence[st class _RawCacheRead(Protocol): - async def async_get_cache(self, *, key: str) -> object: ... + async def async_get_cache(self, *, key: str, local_only: bool = ...) -> object: ... def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead: @@ -910,8 +912,7 @@ def _global_proxy_budget_check(global_proxy_spend: float | None, skip_budget_che and not skip_budget_checks and global_proxy_spend is not None and RouteChecks.is_llm_api_route(route=route) - and route != "/v1/models" - and route != "/models" + and route not in MODEL_LIST_ROUTES ): if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget: raise litellm.BudgetExceededError( @@ -1028,8 +1029,7 @@ async def check_tools_allowlist( # budget cannot reach side-effectful routes like /health/services (Slack/email/webhook). See #27923. MODEL_DISCOVERY_ROUTES: Final = frozenset( { - "/v1/models", - "/models", + *MODEL_LIST_ROUTES, "/model/info", "/v1/model/info", "/v2/model/info", @@ -1800,6 +1800,20 @@ async def get_team_member_default_budget( return budget +def resolve_effective_team_member_budget( + member_budget: LiteLLM_BudgetTable | None, + team_default_budget: LiteLLM_BudgetTable | None, + now: datetime, +) -> float | None: + if member_budget is not None and member_budget.max_budget is not None: + return member_budget.effective_max_budget(now) + if team_default_budget is None or team_default_budget.max_budget is None or team_default_budget.max_budget <= 0: + return None + return team_default_budget.max_budget + ( + member_budget.active_temp_budget_increase(now) if member_budget is not None else 0.0 + ) + + async def resolve_default_end_user_budget( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, @@ -2712,7 +2726,7 @@ async def get_team_membership( use_writer=True, ) - l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) + l1_cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=_key, local_only=True) l1_membership: Final = _membership_from_cached_payload(l1_cached) if not isinstance(l1_membership, _TeamMembershipCacheMiss): return l1_membership @@ -6447,8 +6461,29 @@ def _search_tool_deny_by_default(general_settings: Mapping[str, object]) -> bool return True +def _is_dashboard_session(valid_token: UserAPIKeyAuth | None) -> bool: + return valid_token is not None and valid_token.team_id == UI_TEAM_ID + + +def _is_search_tool_strict_identity(valid_token: UserAPIKeyAuth) -> bool: + if not _is_dashboard_session(valid_token): + return _is_strict_grant_identity(valid_token) + return valid_token.user_role not in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ) + + +def _search_tool_strict_layers( + valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None +) -> frozenset[GrantLayer]: + if _is_dashboard_session(valid_token): + return frozenset(("user",)) + return _strict_grant_layers(True, valid_token, team_object) + + def is_search_tool_deny_by_default_applied(valid_token: UserAPIKeyAuth, general_settings: Mapping[str, object]) -> bool: - return _search_tool_deny_by_default(general_settings) and _is_strict_grant_identity(valid_token) + return _search_tool_deny_by_default(general_settings) and _is_search_tool_strict_identity(valid_token) def _search_tool_denied(object_type: GrantLayer, message: str) -> ProxyException: @@ -6492,12 +6527,12 @@ async def resolve_search_tool_grants( load_team_object: TeamObjectLoader, ) -> SearchToolGrants: """ - Without `general_settings.search_tool_deny_by_default`, or for the master key and dashboard sessions, + Without `general_settings.search_tool_deny_by_default`, or for the master key and admin dashboard sessions, the key and team lists only restrict when nonempty. With it, the identities `_strict_grant_layers` names must each list the tool, and a missing record, `null`, `[]`, or a team or user that fails to load grants nothing """ if not is_search_tool_deny_by_default_applied(valid_token, general_settings): - team_object: Final = await load_team_object() + team_object: Final = await load_team_object() if not _is_dashboard_session(valid_token) else None return SearchToolGrants( grants=MappingProxyType( {"key": valid_token.object_permission, "team": team_object.object_permission if team_object else None} @@ -6507,8 +6542,10 @@ async def resolve_search_tool_grants( from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - strict_team_object: Final = await _strict_team_object(load_team_object) - strict_layers: Final = _strict_grant_layers(True, valid_token, strict_team_object) + strict_team_object: Final = ( + await _strict_team_object(load_team_object) if not _is_dashboard_session(valid_token) else None + ) + strict_layers: Final = _search_tool_strict_layers(valid_token, strict_team_object) if prisma_client is None: return SearchToolGrants(grants=MappingProxyType({}), strict_layers=strict_layers) user_object: Final = ( @@ -6593,7 +6630,7 @@ def check_unregistered_search_fallback( """The unregistered provider fallback names no search tool that a grant could list, so deny-by-default denies it""" if not is_search_tool_deny_by_default_applied(valid_token, general_settings): return True - layer: Final = min(_strict_grant_layers(True, valid_token, None), key=("key", "team", "user").index) + layer: Final = min(_search_tool_strict_layers(valid_token, None), key=("key", "team", "user").index) raise _search_tool_denied( layer, "No registered search tool is available and search_tool_deny_by_default is enabled", @@ -7035,29 +7072,24 @@ async def _check_team_member_budget( # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. - team_member_budget: float | None = None member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None now: Final = get_utc_datetime() - if member_budget_row is not None and member_budget_row.max_budget is not None: - team_member_budget = member_budget_row.effective_max_budget(now=now) - else: - default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") - if isinstance(default_budget_id, str): - default_budget: Final = await get_team_member_default_budget( - budget_id=default_budget_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - # Treat 0 on the team default as "no cap". - # Per-member rows still respect 0 as an explicit admin disable. - if ( - default_budget is not None - and default_budget.max_budget is not None - and default_budget.max_budget > 0 - ): - team_member_budget = default_budget.max_budget + ( - member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 - ) + default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") + team_default_budget: Final = ( + await get_team_member_default_budget( + budget_id=default_budget_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if (member_budget_row is None or member_budget_row.max_budget is None) + and isinstance(default_budget_id, str) + else None + ) + team_member_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget_row, + team_default_budget=team_default_budget, + now=now, + ) if team_member_budget is not None: team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index ae473a62271..2eb25cc6105 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -12,7 +12,7 @@ import litellm from litellm._logging import verbose_proxy_logger, verbose_proxy_stdout_logger from litellm.constants import EMPTY_MAPPING from litellm.integrations.otel.runtime import seed_request_identity -from litellm.litellm_core_utils.core_helpers import is_expected_client_error +from litellm.litellm_core_utils.core_helpers import as_str_mapping, is_expected_client_error from litellm.proxy._types import ( LitellmUserRoles, ModelAccessDeniedProxyException, @@ -41,7 +41,7 @@ DB_UNAVAILABLE_FALLBACK_USER_ID: Final = "__db_unavailable_fallback__" if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any @@ -139,7 +139,7 @@ class UserAPIKeyAuthExceptionHandler: async def handle_authentication_error( e: Exception, request: Request, - request_data: dict[str, object], + request_data: object, route: str, parent_otel_span: Span | None, api_key: str, @@ -190,6 +190,8 @@ class UserAPIKeyAuthExceptionHandler: request_route=route, ) else: + request_mapping: Final = as_str_mapping(request_data) + request_context: Final[dict[str, object]] = dict(request_mapping) if request_mapping is not None else {} # raise the exception to the caller requester_ip: Final = get_request_ip_address( request=request, @@ -229,7 +231,7 @@ class UserAPIKeyAuthExceptionHandler: # without this the failed trace would carry no team/key attributes. seed_request_identity( user_api_key_dict, - model=request_data.get("model"), + model=request_context.get("model"), ) # Budget checks live in tenant-scoped helpers (key / team / org / tag) @@ -242,14 +244,14 @@ class UserAPIKeyAuthExceptionHandler: resolve_llm_provider_for_rate_limit, ) - budget_model: Final = request_data.get("model") + budget_model: Final = request_context.get("model") _, e.llm_provider = resolve_llm_provider_for_rate_limit( budget_model if isinstance(budget_model, str) else None ) # Allow callbacks to transform the error response transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( - request_data=_with_client_context(request_data, requester_ip, user_agent), + request_data=_with_client_context(request_context, requester_ip, user_agent), original_exception=e, user_api_key_dict=user_api_key_dict, error_type=ProxyErrorTypes.auth_error, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 7a186d84aa5..30ede935cfa 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -736,7 +736,9 @@ def is_request_body_safe( def _reject_request_credential_params(litellm_params: Mapping[str, object], *, manages_deployments: bool) -> None: if manages_deployments: - reject_server_owned_wif_fields(litellm_params) + reject_server_owned_wif_fields( + MappingProxyType({key: value for key, value in litellm_params.items() if key != "chatgpt_auth_file"}) + ) else: reject_server_owned_wif_params(litellm_params) reject_federated_credential_reference(litellm_params) @@ -924,8 +926,9 @@ def get_request_route(request: Request) -> str: scope: Final = request.scope if not isinstance(scope, dict): return str(request.url.path) - raw_path: Final[str] = str(scope.get("path", request.url.path)) - root_path: Final[str] = str(scope.get("app_root_path", scope.get("root_path", ""))).rstrip("/") + scope_fields: Final[Mapping[str, object]] = MappingProxyType(scope) + raw_path: Final[str] = str(scope_fields.get("path", request.url.path)) + root_path: Final[str] = str(scope_fields.get("app_root_path", scope_fields.get("root_path", ""))).rstrip("/") if not isinstance(raw_path, str): return str(request.url.path) # Strip root_path only when it matches whole path segments — guarding @@ -1520,8 +1523,15 @@ def get_model_rate_limit_from_metadata( metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit", "model_itpm_limit", "model_otpm_limit"], ) -> dict[str, int] | None: - if getattr(user_api_key_dict, metadata_accessor_key): - return getattr(user_api_key_dict, metadata_accessor_key).get(rate_limit_key) + metadata: Final = ( + user_api_key_dict.team_metadata + if metadata_accessor_key == "team_metadata" + else user_api_key_dict.organization_metadata + if metadata_accessor_key == "organization_metadata" + else user_api_key_dict.project_metadata + ) + if metadata: + return metadata.get(rate_limit_key) return None diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index f1ca2c70f3b..0b349e6a733 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -283,7 +283,7 @@ class IPAddressUtils: @staticmethod def get_mcp_client_ip( request: Request, - general_settings: dict[str, Any] | None = None, + general_settings: dict[str, object] | None = None, ) -> str | None: """ Extract client IP from a FastAPI request for MCP access control. @@ -371,4 +371,4 @@ class IPAddressUtils: return client_ip case _HopCountUnset(): pass - return get_request_ip_address(request, use_x_forwarded_for=use_xff) + return get_request_ip_address(request, use_x_forwarded_for=use_xff is True) diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index 36a569bf18f..d810ba4238b 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -21,6 +21,7 @@ from typing import Final from litellm._logging import verbose_proxy_logger from litellm._version import version from litellm.constants import HIBP_RANGE_API_BASE +from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.types.llms.custom_http import httpxSpecialProvider @@ -143,7 +144,7 @@ async def _is_password_breached(password: str, client: AsyncHTTPHandler) -> bool def is_breach_check_enabled(general_settings: Mapping[str, object]) -> bool: - return general_settings.get("password_policy_check_breached_passwords", True) is not False + return not is_offline_mode() and general_settings.get("password_policy_check_breached_passwords", True) is not False async def is_password_breached( diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py index a9f349d1750..e291c78294c 100644 --- a/litellm/proxy/auth/rds_iam_token.py +++ b/litellm/proxy/auth/rds_iam_token.py @@ -3,7 +3,7 @@ import re from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final +from typing import Final, Protocol import httpx @@ -223,11 +223,15 @@ def rds_region_from_hostname(db_host: str) -> str | None: return match.group("region").lower() +class _RdsAuthTokenClient(Protocol): + def generate_db_auth_token(self, *, DBHostname: str, Port: str, DBUsername: str, Region: str | None) -> str: ... + + def generate_iam_auth_token( db_host: str, db_port: str, db_user: str, - client: Any | None = None, + client: _RdsAuthTokenClient | None = None, *, region: str | None = None, ) -> str: diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 720bfabb95a..c402f111563 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -7,6 +7,7 @@ from fastapi import HTTPException, Request, status from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( + MODEL_LIST_ROUTES, UI_TEAM_ID, CommonProxyErrors, KeyManagementRoutes, @@ -93,8 +94,7 @@ _SELF_SCOPED_CREDENTIAL_ROUTES: Final = ( _UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( { - "/models", - "/v1/models", + *MODEL_LIST_ROUTES, "/utils/token_counter", "/utils/transform_request", "/responses/input_tokens", @@ -160,7 +160,7 @@ class RouteChecks: if allowed_route in LiteLLMRoutes._member_names_: if RouteChecks.check_route_access( route=route, - allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value, + allowed_routes=LiteLLMRoutes[allowed_route].value, ): if ( allowed_route in _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 27fa4731baf..b1e8bdbe787 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -13,6 +13,7 @@ import re import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import Final, NamedTuple, Protocol, Union, cast import fastapi @@ -177,6 +178,7 @@ from litellm.proxy.utils import ( normalize_route_for_root_path, ) from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository +from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model from litellm.router_utils.fallback_event_handlers import ( @@ -1319,7 +1321,7 @@ async def fetch_global_spend_with_event_coordination( async def _load_global_spend() -> float | None: proxy_budget_row: Final = await bounded_db_lookup( - prisma_client.db.litellm_usertable.find_unique(where={"user_id": LITELLM_PROXY_BUDGET_NAME}), + UserRepository(prisma_client).table.find_unique(where={"user_id": LITELLM_PROXY_BUDGET_NAME}), name="proxy_budget", ) return float(proxy_budget_row.spend) if proxy_budget_row is not None else None @@ -1635,7 +1637,7 @@ async def _auto_register_jwt_mapping( if minted: try: async with db_span("delete_orphaned_jwt_key", "LiteLLM_VerificationToken"): - await prisma_client.db.litellm_verificationtoken.delete(where={"token": token_hash}) + await VerificationTokenRepository(prisma_client).table.delete(where={"token": token_hash}) except Exception as delete_err: # Don't fail the request if cleanup fails — the orphan is # unmapped and inert. Log so an operator can prune it later. @@ -2693,7 +2695,7 @@ async def user_api_key_auth_builder( and valid_token.last_refreshed_at is not None and team_obj.last_refreshed_at > valid_token.last_refreshed_at ): - team_obj_dict: Final = team_obj.__dict__ + team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__) for k, v in team_obj_dict.items(): field_name = f"team_{k}" diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index 597080be3b3..775d00ed53e 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request @@ -19,6 +20,10 @@ router: Final = APIRouter( ) +def _copy_cache_params(cache_params: Mapping[str, object]) -> Mapping[str, object]: + return dict(**cache_params) + + def _extract_cache_params() -> dict[str, object]: """ Safely extracts and cleans cache parameters. @@ -41,7 +46,7 @@ def _extract_cache_params() -> dict[str, object]: try: cache_params: Final = vars(litellm.cache.cache) cleaned_params: Final = ( - HealthCheckCacheParams.model_validate(dict(**cache_params)).model_dump() if cache_params else {} + HealthCheckCacheParams.model_validate(_copy_cache_params(cache_params)).model_dump() if cache_params else {} ) return masker.mask_dict(cleaned_params) except (AttributeError, TypeError) as e: diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index c027bf2fe98..ccdef44f9e8 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -59,7 +59,7 @@ class ChatClient: frequency_penalty: float | None = None, user: str | None = None, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Create a chat completion. @@ -184,7 +184,8 @@ class ChatClient: response.raise_for_status() # Parse SSE stream - for line in response.iter_lines(): + lines: Final[Iterator[object]] = response.iter_lines() + for line in lines: if line: line = _SSE_LINE.validate_python(line).decode("utf-8") if line.startswith("data: "): diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 842705ad5a7..4038ff6f582 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -494,7 +494,7 @@ Each command resolves your LiteLLM key (logging in via SSO when none is stored a The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. Codex gets the same list as a catalog file, `$CODEX_HOME/litellm-models.json` (default `~/.codex/`), passed as `-c model_catalog_json=` so `/model` lists exactly the proxy's chat models; a proxy model the installed Codex already knows (`gpt-5.5`, say) keeps that Codex's own entry, reasoning levels and prompt included, and only its place in the picker comes from the proxy, while a model Codex does not know gets the plain entry Codex uses for an unknown `-m` slug. Before launching, `lite codex` asks the installed Codex for its own list and then has it read the written file back (both through `codex debug models`), and when the fetch, either of those or the write fails (Codex releases older than 0.130 have no such command) it says so on stderr and launches with Codex's built-in catalog, leaving a rejected file in place. -pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. +pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. The entry also turns on pi's session header (`x-session-id`, sent by pi 0.80.7 and later), so every chat turn of one pi session shares one session id in the logs and sticks to one deployment when session affinity is on. A custom `pi --session-id` must be at least 8 letters, digits, `_`, or `-` for the proxy to read it. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. Options (these belong to the wrapper, so put them before the agent's own flags): diff --git a/litellm/proxy/client/cli/commands/encryption.py b/litellm/proxy/client/cli/commands/encryption.py index ef73b4f0d2d..61c8db617af 100644 --- a/litellm/proxy/client/cli/commands/encryption.py +++ b/litellm/proxy/client/cli/commands/encryption.py @@ -4,9 +4,14 @@ from typing import Final import click import rich +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import as_str_mapping from ...http_client import HTTPClient -from .auth import cli_context_obj +from ._cli_context import cli_context_values + +_DECODED_JSON: Final = TypeAdapter(object) @click.group() @@ -40,23 +45,32 @@ def migrate(ctx: click.Context, check_only: bool, dry_run: bool): lite encryption migrate --check # attestation scan, no writes lite encryption migrate # perform the migration """ - client: Final = HTTPClient(cli_context_obj(ctx)["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = HTTPClient(context["base_url"], context["api_key"]) - if check_only: - response = client.request("GET", "/credentials/migrate-encryption/check") - else: - response = client.request( + response: Final = _DECODED_JSON.validate_python( + client.request("GET", "/credentials/migrate-encryption/check") + if check_only + else client.request( "POST", "/credentials/migrate-encryption", json={}, params={"dry_run": "true"} if dry_run else None, ) + ) rich.print_json(data=response) - report: Final = response.get("report", {}) if isinstance(response, dict) else {} + report_value: Final = (as_str_mapping(response) or {}).get("report", {}) + report: Final = as_str_mapping(report_value) + if report is None: + raise click.ClickException(f"Unexpected report in the proxy response: {report_value!r}") residual: Final = report.get("residual_legacy") - if residual is not None and residual > 0: + if residual is None: + return + if not isinstance(residual, int | float): + raise click.ClickException(f"Unexpected residual_legacy in the proxy response: {residual!r}") + if residual > 0: rich.print(f"[yellow]Residual legacy values remaining: {residual}[/yellow]") elif residual == 0: rich.print("[green]No legacy values remaining (residual_legacy == 0).[/green]") diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index b228617699c..357b82bdf67 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -191,7 +191,7 @@ def provider_block( Real contextWindow/maxTokens matter: pi otherwise assumes 128k/16384, which breaks compaction thresholds and over-asks models with smaller output caps. pi sniffs compat from the base URL, and one gateway URL fronts models with - different capabilities, so both flags are pinned off. + different capabilities, so the store and long-cache flags are pinned off. """ return { "baseUrl": base_url.rstrip("/") + "/v1", @@ -199,6 +199,8 @@ def provider_block( "compat": { "supportsStore": False, "supportsLongCacheRetention": False, + "sendSessionAffinityHeaders": True, + "sessionAffinityFormat": "openrouter", }, "apiKey": f"${LITELLM_PROXY_API_KEY_ENV}", "models": [_model_entry(model_id, limits) for model_id in model_ids], diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 5aeefa0c739..fef3e313ee5 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -89,6 +89,28 @@ class Fetch(Protocol): def __call__(self, credentials: Credentials, session_id: str) -> Fetched: ... +class _ResponseBody(Protocol): + def read(self) -> bytes: ... + + +class _OpenedResponse(Protocol): + def __enter__(self) -> _ResponseBody: ... + + def __exit__(self, *exc_info: object) -> None: ... + + +class _UrlOpener(Protocol): + def __call__(self, url: urllib.request.Request, /, *, timeout: float) -> _OpenedResponse: ... + + +def _url_opener() -> _UrlOpener: + return urllib.request.urlopen + + +def _json_decoder() -> Callable[[bytes | str], object]: + return json.loads + + def as_mapping(value: object) -> Mapping[str, object]: return value if isinstance(value, dict) else EMPTY @@ -106,7 +128,7 @@ def printable(value: object) -> str: def load_json(raw: bytes | str) -> object: try: - return json.loads(raw) + return _json_decoder()(raw) except ValueError: return None @@ -195,7 +217,7 @@ def fetch_session(credentials: Credentials, session_id: str) -> Fetched: }, ) try: - with urllib.request.urlopen(request, timeout=FETCH_TIMEOUT_SECONDS) as response: + with _url_opener()(request, timeout=FETCH_TIMEOUT_SECONDS) as response: raw: Final[bytes] = response.read() except urllib.error.HTTPError as error: return Fetched(session=None, definitive=400 <= error.code < 500) diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index f7c545ef48e..7c01732119e 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -189,7 +189,7 @@ def execute_command(path: ShellScope, ctx: click.Context, cli: click.Group) -> N try: if path[0].startswith("-"): cli.main( - ["--base-url", ctx.obj["base_url"], *path], + ["--base-url", cli_context_obj(ctx)["base_url"], *path], prog_name="lite", obj={"secret_vault": context_secret_vault(ctx)}, standalone_mode=False, diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 3f1e7f4259c..a1964d8a1bb 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -4,6 +4,7 @@ from collections.abc import Callable, Iterator, Mapping, MutableMapping from typing import Final, Literal, Protocol import click +from pydantic import TypeAdapter from litellm._version import version as litellm_version from litellm.litellm_core_utils.cli_token_utils import load_cli_token @@ -89,10 +90,13 @@ class LazyCommands(MutableMapping[str, click.Command]): return len(self._loaders) +_MODULE_ATTRIBUTE: Final = TypeAdapter(object) + + def _command_in(module: str, attribute: str) -> Callable[[], click.Command]: @functools.cache def load() -> click.Command: - match getattr(importlib.import_module(module, __package__), attribute): + match _MODULE_ATTRIBUTE.validate_python(getattr(importlib.import_module(module, __package__), attribute)): case click.Command() as command: return command case _: diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index ebcd79457a6..649fe529936 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -75,7 +75,7 @@ class CredentialsManagementClient: credential_values: Mapping[str, object], return_request: bool = False, display_name: str | None = None, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Create a new credential. @@ -121,7 +121,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Delete a credential by name. @@ -182,7 +182,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Get a credential by name. diff --git a/litellm/proxy/client/http_client.py b/litellm/proxy/client/http_client.py index aa0b986b1ad..96a170c75cb 100644 --- a/litellm/proxy/client/http_client.py +++ b/litellm/proxy/client/http_client.py @@ -4,6 +4,11 @@ from collections.abc import Mapping from typing import Any, Final import requests +from typing_extensions import TypedDict, Unpack + + +class _RequestOptions(TypedDict, total=False, extra_items=object): + pass class HTTPClient: @@ -29,7 +34,7 @@ class HTTPClient: data: Mapping[str, object] | list | bytes | None = None, json: Mapping[str, object] | list | None = None, headers: dict[str, str] | None = None, - **kwargs: Any, + **kwargs: Unpack[_RequestOptions], ) -> Any: """Make an HTTP request to the LiteLLM proxy server. diff --git a/litellm/proxy/client/model_groups.py b/litellm/proxy/client/model_groups.py index fef307600c4..1555d170212 100644 --- a/litellm/proxy/client/model_groups.py +++ b/litellm/proxy/client/model_groups.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final import requests @@ -31,7 +31,7 @@ class ModelGroupsManagementClient: headers["Authorization"] = f"Bearer {self._api_key}" return headers - def info(self, return_request: bool = False) -> list[dict[str, Any]] | requests.Request: + def info(self, return_request: bool = False) -> list[dict[str, object]] | requests.Request: """ Get detailed information about all model groups from the server. diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 1b50f8c24f3..1361a3e6364 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -158,7 +158,7 @@ class ModelsManagementClient: model_id: str | None = None, model_name: str | None = None, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Get information about a specific model by its ID or name. @@ -168,7 +168,7 @@ class ModelsManagementClient: return_request (bool): If True, returns the prepared request object instead of executing it Returns: - Union[Dict[str, Any], requests.Request]: Either the model information from the server or + dict[str, object] | requests.Request: Either the model information from the server or a prepared request object if return_request is True Raises: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 59db23820fc..61cf4d696a1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -67,6 +67,7 @@ from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.litellm_core_utils.llm_response_utils.get_headers import ( @@ -118,6 +119,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( ) from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression +from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( route_request, @@ -265,6 +267,7 @@ from litellm.proxy.litellm_pre_call_utils import ( sanitize_for_log, ) from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -551,22 +554,29 @@ def model_deployments_support_stream_options( return all(_litellm_model_supports_stream_options(m) for m in candidate_models) +_STREAM_USAGE_FLAGS: Final = ("_litellm_strip_stream_usage", STREAM_USAGE_INJECTED_FIELD) + + def stream_usage_tracking_updates( data: Mapping[str, object], general_settings: Mapping[str, object], route_type: str, supports_stream_options: Callable[[], bool], ) -> Mapping[str, object]: - scrub: Final = {"_litellm_strip_stream_usage": False} if "_litellm_strip_stream_usage" in data else {} + scrub: Final = {flag: False for flag in _STREAM_USAGE_FLAGS if flag in data} if data.get("stream", False) is not True: return scrub always_include: Final = general_settings.get("always_include_stream_usage") stream_options: Final = data.get("stream_options") if always_include is True: if "stream_options" not in data: - return {**scrub, "stream_options": {"include_usage": True}} + return {**scrub, "stream_options": {"include_usage": True}, STREAM_USAGE_INJECTED_FIELD: True} if isinstance(stream_options, dict) and "include_usage" not in stream_options: - return {**scrub, "stream_options": {**stream_options, "include_usage": True}} + return { + **scrub, + "stream_options": {**stream_options, "include_usage": True}, + STREAM_USAGE_INJECTED_FIELD: True, + } return scrub if always_include is False and route_type == "agenerate_content_stream" and "stream_options" not in data: return {**scrub, "stream_options": {"include_usage": False}} @@ -580,6 +590,7 @@ def stream_usage_tracking_updates( return { "stream_options": {**merged_stream_options, "include_usage": True}, "_litellm_strip_stream_usage": True, + STREAM_USAGE_INJECTED_FIELD: True, } @@ -1406,6 +1417,8 @@ async def open_sse_before_first_byte( if interval is None: return await produce_response + # Share admission state with stream cleanup before the producer copies this context. + get_or_create_request_stash() produce_task: Final = asyncio.ensure_future(produce_response) await asyncio.wait((produce_task,), timeout=interval) if produce_task.done(): @@ -1975,6 +1988,7 @@ class ProxyBaseLLMRequestProcessing: self.data = data self.served_model = served_model self._tags_before_guardrails: frozenset[str] | None = None + self.route_type: str | None = None @property def litellm_call_id(self) -> str | None: @@ -2424,6 +2438,7 @@ class ProxyBaseLLMRequestProcessing: ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call ## IMPORTANT Note: - initialize this before running pre-call checks. Ensures we log rejected requests to langfuse. + self.route_type = route_type logging_obj, self.data = litellm.utils.function_setup( original_function=route_type, rules_obj=litellm.utils.Rules(), @@ -2790,6 +2805,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=self.data, + dispatched_call_type=self.route_type, ) return await open_sse_before_first_byte( @@ -2854,6 +2870,7 @@ class ProxyBaseLLMRequestProcessing: "skip_pre_call_logic=True requires litellm_logging_obj to be set in data. " "Ensure common_processing_pre_call_logic was called before using this parameter." ) + self.route_type = route_type else: self.data, logging_obj = await self._pre_call_with_fallbacks( request=request, @@ -3463,6 +3480,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=request_data, + dispatched_call_type=self.route_type, ) except HTTPException: return @@ -3987,6 +4005,8 @@ class ProxyBaseLLMRequestProcessing: async def _on_deferred_native_stream_complete(*args: object) -> None: match args: + case (BilledRoundsDispatch(dispatch=logging_coroutine),): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (assembled_response, cache_hit): @@ -4130,6 +4150,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=e, request_data=self.data, + dispatched_call_type=self.route_type, ) # Use transformed exception if callback returned one, otherwise use original if transformed_exception is not None: @@ -4803,10 +4824,15 @@ class ProxyBaseLLMRequestProcessing: model_info = metadata.get("model_info") or {} model_id = model_info.get("id", None) - # 3. Final fallback to self.data["litellm_metadata"] (for routes like /v1/responses that populate data before error) + # 3. Final fallback to the request metadata the router updates in place with the selected deployment if not model_id: - litellm_metadata: Final = self.data.get("litellm_metadata", {}) or {} - model_info = litellm_metadata.get("model_info", {}) or {} - model_id = model_info.get("id", None) + model_id = self._metadata_model_id("litellm_metadata") or self._metadata_model_id("metadata") return model_id + + def _metadata_model_id(self, metadata_key: str) -> str | None: + match self.data.get(metadata_key): + case {"model_info": {"id": str() as request_model_id}}: + return request_model_id + case _: + return None diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 40cdb117476..b2640d2a765 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,11 +1,15 @@ import asyncio import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter + from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # legacy module exports ConfigSyncPubSub, _ConfigSyncPubSub, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -30,6 +34,7 @@ _pending_publishes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong _in_flight_publishes: Final = asyncio.Semaphore(_MAX_IN_FLIGHT_PUBLISHES) _BACKOFF_INITIAL_SECONDS: Final = 5.0 _BACKOFF_MAX_SECONDS: Final = 60.0 +_DECODED_JSON: Final = TypeAdapter(object) def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @@ -46,7 +51,9 @@ class _CacheInvalidationMessage: def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str: - message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) + message: Final[Mapping[str, object]] = MappingProxyType( + asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) + ) return json.dumps({field: value for field, value in message.items() if value is not None}) @@ -62,10 +69,10 @@ def _message_from_data(data: object) -> _CacheInvalidationMessage | None: if not isinstance(data, str): return None try: - parsed: Final = json.loads(data) + parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(data))) except json.JSONDecodeError: return None - if not isinstance(parsed, dict): + if parsed is None: return None cache_key: Final = parsed.get("cache_key") if not isinstance(cache_key, str): diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 46434a997b7..f889748c504 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from itertools import accumulate from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict, assert_never import litellm @@ -39,6 +40,7 @@ from litellm.types.utils import ( ) _CALLBACK_VAR_MASKER: Final = SensitiveDataMasker() +_LOADED_CALLBACK: Final = TypeAdapter(object) # Compound names that are credential-bearing but don't contain any of the # default sensitive segments (so SensitiveDataMasker won't flag them). _EXTRA_SENSITIVE_CALLBACK_KEYS: Final = {"gcs_path_service_account"} @@ -389,9 +391,11 @@ def initialize_callbacks_on_proxy( imported_list.append( _loaded_callback_or_raise( entry=callback, - loaded=get_instance_fn( - value=callback, - config_file_path=config_file_path, + loaded=_LOADED_CALLBACK.validate_python( + get_instance_fn( + value=callback, + config_file_path=config_file_path, + ) ), ) ) @@ -408,9 +412,11 @@ def initialize_callbacks_on_proxy( litellm.callbacks = [ _loaded_callback_or_raise( entry=value, - loaded=get_instance_fn( - value=value, - config_file_path=config_file_path, + loaded=_LOADED_CALLBACK.validate_python( + get_instance_fn( + value=value, + config_file_path=config_file_path, + ) ), ) ] diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 33633150711..ec4baa76514 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -11,12 +11,14 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import TypeAdapter from typing_extensions import ReadOnly from litellm import get_secret_str from litellm._logging import verbose_proxy_logger from litellm.constants import PYTHON_GC_THRESHOLD from litellm.litellm_core_utils.bug_report import EnvironmentReport +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.bug_report_config import build_proxy_environment_report @@ -26,6 +28,7 @@ if TYPE_CHECKING: from litellm.router import Router router: Final = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) # Configure garbage collection thresholds from environment variables @@ -885,8 +888,8 @@ def init_verbose_loggers(): return if os.path.isfile(worker_config): return - _settings: Final = json.loads(worker_config) - if not isinstance(_settings, dict): + _settings: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(worker_config))) + if _settings is None: return debug: Final = _settings.get("debug", None) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 245a56678eb..9f7fed40ab3 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -6,7 +6,7 @@ from typing import Annotated, Any, Final, Literal, Union, get_args, get_origin import orjson from fastapi import Request, UploadFile, status -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from starlette._utils import get_route_path from starlette.requests import HTTPConnection from typing_extensions import NotRequired, ReadOnly, Required, assert_never @@ -31,6 +31,7 @@ _BINARY_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-protobu _ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required}) _DECODED_JSON: Final = TypeAdapter(object) +_JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) def resolve_inference_model( @@ -161,13 +162,11 @@ def coerce_numeric_form_fields( } -def _parse_binary_body(body: bytes) -> dict: +def _parse_binary_body(body: bytes) -> dict[str, object]: """JSON sent under a binary content type still parses; real binary (protobuf) carries no params -> {}.""" try: - parsed: Final = orjson.loads(body) - if isinstance(parsed, dict): - return parsed - except orjson.JSONDecodeError: + return dict(_JSON_REQUEST_BODY.validate_python(orjson.loads(body))) + except (orjson.JSONDecodeError, ValidationError): pass return {} @@ -272,7 +271,7 @@ async def read_request_body(request: Request | None) -> dict: parsed_body = {} else: try: - parsed_body = orjson.loads(body) + parsed_body = dict(_JSON_REQUEST_BODY.validate_python(orjson.loads(body))) except orjson.JSONDecodeError as e: # The surrogate-repair fallback below runs two full-body re.sub # passes, which block the event loop on multi-MB malformed bodies. @@ -298,8 +297,9 @@ async def read_request_body(request: Request | None) -> dict: body_str = re.sub(r"(? dict: safe_set_request_parsed_body(request=request, parsed_body=parsed_body) return parsed_body + except ValidationError as e: + from fastapi.routing import APIRoute + + from litellm.proxy.auth.auth_utils import ( + request_dispatched_to_pass_through_endpoint, + request_dispatched_to_provider_pass_through, + ) + from litellm.proxy.lens.buffering import LensRoute + + route: Final[object] = request.scope.get("route") if request is not None else None + if request is not None and ( + request_dispatched_to_pass_through_endpoint(request) + or request_dispatched_to_provider_pass_through(request) + or isinstance(route, LensRoute) + or ( + isinstance(route, APIRoute) + and route.body_field is not None + and not route.body_field.field_info.is_required() + and (await request.body()).strip() == b"null" + ) + ): + return {} + raise ProxyException( + message="Invalid JSON payload: expected an object", + type="invalid_request_error", + param="request_body", + code=status.HTTP_400_BAD_REQUEST, + ) from e except (json.JSONDecodeError, orjson.JSONDecodeError, ProxyException) as e: # Re-raise ProxyException as-is verbose_proxy_logger.error("Invalid JSON payload received: %s", e) @@ -415,7 +443,9 @@ def rewrite_request_model( if request is None: return cached_body: Final = _safe_get_request_parsed_body(request=request) - body: Final = {**cached_body, "model": model} if cached_body is not None else request_data + body: Final = ( + {**_JSON_REQUEST_BODY.validate_python(cached_body), "model": model} if cached_body is not None else request_data + ) safe_set_request_parsed_body(request=request, parsed_body=body) request._json = body # pyright: ignore[reportPrivateUsage] # Starlette JSON cache request._body = orjson.dumps(body) # pyright: ignore[reportPrivateUsage] # Starlette body cache diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 4a082eb307b..640a947f9c5 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase _BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_YAML: Final = TypeAdapter(object) class BucketObjectFetcher(Protocol): @@ -31,7 +32,7 @@ class SyncBucketObjectReader(Protocol): def _parsed_config(object_key: str, file_contents: str) -> object | None: try: - parsed: Final = yaml.safe_load(file_contents) + parsed: Final = _DECODED_YAML.validate_python(yaml.safe_load(file_contents)) except yaml.YAMLError as e: verbose_proxy_logger.error("Config object %s is not valid YAML: %s", object_key, e) return None diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index 089d0b7a226..02da6f706d1 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -37,7 +37,7 @@ This module provides a single proxy-side error class that: import json from collections.abc import Mapping -from typing import Any, Final +from typing import Final from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError @@ -171,7 +171,7 @@ class ProxyRateLimitError(HTTPException, RateLimitError): def __init__( self, - detail: Any, + detail: object, headers: Mapping[str, object] | None = None, category: str | RateLimitErrorCategory = RateLimitErrorCategory.LITELLM_RATE_LIMIT, rate_limit_type: str | RateLimitType | None = None, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index ccf2232e379..59fd00cec43 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -65,6 +65,7 @@ from litellm.repositories.table_repositories import ( from litellm.repositories.team_repository import TeamRepository from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, + SpendResetWrites, budget_cascade_unit_of_work, enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, @@ -235,7 +236,7 @@ def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]: def _queue_budget_linked_resets( - writes: LinkedSpendResetWrites, + writes: SpendResetWrites, cascade: "_BudgetCascade", extra: Mapping[str, object] = MappingProxyType({}), ) -> None: @@ -908,6 +909,7 @@ class ResetBudgetJob: db_span("reset_budget_cascade", "LiteLLM_BudgetTable"), budget_cascade_unit_of_work(self._new_batch) as uow, ): + uow.team_memberships.queue_total_spend_repair(cascade.budget_ids) _queue_budget_linked_resets(uow.team_memberships, cascade) _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 35bd69ef0bd..a2ffdd4d59d 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -13,6 +13,7 @@ from typing import Final from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -26,6 +27,8 @@ from litellm.proxy.container_endpoints.ownership import ( get_container_forwarding_params, ) +_UPLOADED_FILE_FIELD: Final = TypeAdapter(object) + def _load_endpoints_config() -> dict: """Load the endpoints configuration from JSON file.""" @@ -304,7 +307,7 @@ async def _process_multipart_upload_request( raise HTTPException(status_code=400, detail="Missing required 'file' field") # convert_upload_files_to_file_data returns list of tuples, extract single file - file_list: Final = data["file"] + file_list: Final = _UPLOADED_FILE_FIELD.validate_python(data["file"]) if isinstance(file_list, list) and len(file_list) > 0: data["file"] = file_list[0] diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index aeb0bb1e3d0..9e4ab3c3c4e 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,13 +1,14 @@ import json from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet -from typing import TYPE_CHECKING, Any, Final, TypeAlias +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias from fastapi import HTTPException from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -111,12 +112,20 @@ def _get_response_id(response: object) -> str | None: return value if isinstance(value, str) else None -def _dump_response(response: Any) -> dict[str, object]: +class _SupportsModelDump(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _SupportsPydanticDict(Protocol): + def dict(self) -> Mapping[str, object]: ... + + +def _dump_response(response: object) -> Mapping[str, object]: if isinstance(response, dict): return dict(response) - if hasattr(response, "model_dump"): + if has_attributes(response, _SupportsModelDump, "model_dump"): return response.model_dump() - if hasattr(response, "dict"): + if has_attributes(response, _SupportsPydanticDict, "dict"): return response.dict() return {"id": _get_response_id(response)} @@ -199,9 +208,11 @@ async def record_container_owner( original_container_id, resolved_provider = decode_container_id_for_ownership(container_id, custom_llm_provider) model_object_id: Final = _container_model_object_id(original_container_id, resolved_provider) - file_object: Final = _dump_response(response) - file_object["custom_llm_provider"] = resolved_provider - file_object["provider_container_id"] = original_container_id + file_object: Final = { + **_dump_response(response), + "custom_llm_provider": resolved_provider, + "provider_container_id": original_container_id, + } # Prisma Python requires Json fields to be serialized as a JSON string. file_object_json: Final[str] = json.dumps(file_object) diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 75882bc3fb1..4ddcf7e600b 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -526,7 +526,7 @@ async def _upsert_turn_with_retry( except DB_RETRY_SAFE_ERROR_TYPES: if attempt >= n_retry_times: raise - await asyncio.sleep(2**attempt + random.uniform(0, 1)) + await asyncio.sleep(2.0**attempt + random.uniform(0, 1)) else: return diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4e98823ba77..8ac2597c99b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,11 +12,11 @@ import json import os import random import time -import traceback from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable, Mapping, Sequence from contextlib import asynccontextmanager from contextvars import ContextVar from datetime import datetime, timedelta, timezone +from functools import reduce from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload from urllib.parse import quote, unquote @@ -37,7 +37,6 @@ from litellm.constants import ( from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, @@ -74,7 +73,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, ) -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_failure_is_requeue_safe from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.request_error_tracking import request_error_accumulator from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout @@ -90,6 +89,7 @@ from litellm.proxy.spend_tracking.savings import ( ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.repositories.prisma_protocols import BatchTable +from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import CallTypes if TYPE_CHECKING: @@ -104,6 +104,7 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) _SPEND_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_CALL_KWARG_VALUE: Final = TypeAdapter(object) _JsonPayload: TypeAlias = Mapping[str, object] | Sequence[object] | str | int | float | bool | None _JSON_VALUE_ADAPTER: Final[TypeAdapter[_JsonPayload]] = TypeAdapter( _JsonPayload, config=ConfigDict(strict=True, hide_input_in_errors=True) @@ -249,16 +250,6 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]): ) -> None: ... -_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) - - -def _spend_commit_failure_is_requeue_safe(e: Exception) -> bool: - if isinstance(e, DB_CONNECTION_ERROR_TYPES): - return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) - sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) - return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES - - _SpendTableName = Literal[ "user_list_transactions", "end_user_list_transactions", @@ -296,7 +287,7 @@ def _spend_tables_left_to_send( _SPEND_TABLE_COMMIT_ORDER[len(committed)] if len(committed) < len(_SPEND_TABLE_COMMIT_ORDER) else None ) dropped: Final[frozenset[_SpendTableName]] = ( - frozenset() if in_flight is None or _spend_commit_failure_is_requeue_safe(failure) else frozenset({in_flight}) + frozenset() if in_flight is None or db_write_failure_is_requeue_safe(failure) else frozenset({in_flight}) ) if dropped and in_flight is not None: spend_log_error( @@ -313,6 +304,52 @@ def _spend_tables_left_to_send( return cast(DBSpendUpdateTransactions, remaining) if any(remaining.values()) else None +class _SpendRestore(TypedDict, total=False): + db_spend_update_transactions: ReadOnly[DBSpendUpdateTransactions] + daily_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + daily_team_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + daily_org_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + daily_end_user_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + daily_agent_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + window_spend_update_transactions: ReadOnly[Sequence[WindowSpendTransaction]] + + +def _merged_spend_restore(first: _SpendRestore, second: _SpendRestore) -> _SpendRestore: + return {**first, **second} + + +def _report_spend_update_step_failure(step_name: str, *, payload: Mapping[str, object], exc: Exception) -> None: + spend_log_error( + "Spend tracking - %s failed, this request is missing from that spend rollup. request_id=%s api_key=%s", + step_name, + payload.get("request_id"), + payload.get("api_key"), + exc=exc, + ) + + +_EMPTY_SPEND_LOG_METADATA: Final = cast(SpendLogsMetadata, MappingProxyType({})) # cast-ok: empty read-only metadata + + +def _parse_spend_log_metadata(payload: Mapping[str, object]) -> SpendLogsMetadata: + """Malformed metadata costs the request its metadata-derived fields, never its daily rollup row.""" + raw: Final = payload.get("metadata") + if not isinstance(raw, str): + return _EMPTY_SPEND_LOG_METADATA + try: + parsed: Final[SpendLogsMetadata] = json.loads(raw) + except json.JSONDecodeError as e: + spend_log_error( + "Spend tracking - unreadable spend log metadata, rollups for this request are written without " + "metadata-derived fields. request_id=%s api_key=%s", + payload.get("request_id"), + payload.get("api_key"), + exc=e, + ) + return _EMPTY_SPEND_LOG_METADATA + return parsed + + def _is_billed_retry_attempt(metadata: SpendLogsMetadata) -> bool: """An earlier billed retry attempt's own spend row, not a request the caller sent.""" return BILLED_RETRY_ATTEMPT_METADATA_KEY in metadata and metadata[BILLED_RETRY_ATTEMPT_METADATA_KEY] is True @@ -435,6 +472,8 @@ _TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS # recreates the member. The rows travel as one JSON document, not as a numeric array: Prisma # types a raw array parameter from the first batch a connection sees, so after an all-$0 batch # (integers) every later fractional batch on that connection failed with "improper binary format". +# total_spend is floored at spend before adding: rows written before total_spend existed (migration +# 20260421135425, default 0) carry their history only in spend, so a plain increment keeps them behind forever. _TEAM_MEMBER_SPEND_SQL: Final = """ INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend) SELECT member.user_id, member.team_id, member.cost, member.cost @@ -450,7 +489,7 @@ WHERE EXISTS ( ) ON CONFLICT (user_id, team_id) DO UPDATE SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, - total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend + total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend) + EXCLUDED.total_spend """ @@ -647,8 +686,8 @@ class DBSpendUpdateWriter: "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." ) - # Single task replaces 11 create_task() calls - asyncio.create_task( + spawn_logging_task( + asyncio.get_running_loop(), self._batch_database_updates( response_cost=response_cost, user_id=user_id, @@ -661,7 +700,8 @@ class DBSpendUpdateWriter: litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, request_model_access_groups=get_request_model_access_groups(kwargs), - ) + ), + task_name="spend rollups", ) self._enqueue_tool_registry_upsert( @@ -812,7 +852,7 @@ class DBSpendUpdateWriter: spend=payload["spend"], total_tokens=payload["total_tokens"], completion_response=completion_response, - realtime_tool_calls=(kwargs or {}).get("realtime_tool_calls"), + realtime_tool_calls=_CALL_KWARG_VALUE.validate_python((kwargs or {}).get("realtime_tool_calls")), ) if transaction is None: return @@ -1086,91 +1126,52 @@ class DBSpendUpdateWriter: project_id: str | None = None, ): """ - Runs all 13 spend-update helpers sequentially inside a single asyncio task. - - Each helper is wrapped in try/except so one failure doesn't prevent the others. - - The deepcopy runs here, off the awaited request path, so the daily spend - helpers get a payload isolated from the spend-log queue entry and the caller. + Runs every spend-update helper for one request. Each helper goes through + _run_spend_update_step so one failure never skips the sibling writes. The deepcopy + isolates the daily spend helpers' payload from the spend-log queue entry and the caller. """ payload_copy: Final = copy.deepcopy(payload) request_tags: Final = payload_copy.get("request_tags") - try: - await self._update_user_db( + await self._run_spend_update_step( + "_update_user_db", + self._update_user_db( response_cost=response_cost, user_id=user_id, prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, end_user_id=end_user_id, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_user_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_key_db( - response_cost=response_cost, - hashed_token=hashed_token, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_key_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_team_db( - response_cost=response_cost, - team_id=team_id, - user_id=user_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_team_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_org_db( - response_cost=response_cost, - org_id=org_id, - user_id=user_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_org_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_project_db( - response_cost=response_cost, - project_id=project_id, - prisma_client=prisma_client, - ) - except Exception: # noqa: BLE001 # a project enqueue failure must not skip the sibling spend writes - verbose_proxy_logger.debug( - "_batch_database_updates: _update_project_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_tag_db( - response_cost=response_cost, - request_tags=request_tags, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_tag_db failed: %s", - traceback.format_exc(), - ) - + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_key_db", + self._update_key_db(response_cost=response_cost, hashed_token=hashed_token, prisma_client=prisma_client), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_team_db", + self._update_team_db( + response_cost=response_cost, team_id=team_id, user_id=user_id, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_org_db", + self._update_org_db( + response_cost=response_cost, org_id=org_id, user_id=user_id, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_project_db", + self._update_project_db(response_cost=response_cost, project_id=project_id, prisma_client=prisma_client), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_tag_db", + self._update_tag_db(response_cost=response_cost, request_tags=request_tags, prisma_client=prisma_client), + payload=payload_copy, + ) await self._update_model_access_group_db( response_cost=response_cost, request_model_access_groups=request_model_access_groups, @@ -1178,96 +1179,64 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, router=get_llm_router(), ) - - _agent_id_for_spend: Final = payload_copy.get("agent_id") - try: - await self._update_agent_db( - response_cost=response_cost, - agent_id=_agent_id_for_spend, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_agent_db failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_user_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_user_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_end_user_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_end_user_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_agent_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_agent_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_team_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_team_transaction failed: %s", - traceback.format_exc(), - ) + await self._run_spend_update_step( + "_update_agent_db", + self._update_agent_db( + response_cost=response_cost, agent_id=payload_copy.get("agent_id"), prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_user_transaction", + self.add_spend_log_transaction_to_daily_user_transaction(payload=payload_copy, prisma_client=prisma_client), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_end_user_transaction", + self.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload_copy, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_agent_transaction", + self.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload_copy, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_team_transaction", + self.add_spend_log_transaction_to_daily_team_transaction(payload=payload_copy, prisma_client=prisma_client), + payload=payload_copy, + ) try: self._record_request_error(payload=payload_copy, prisma_client=prisma_client) - except Exception: # noqa: BLE001 # the rollup must never skip the sibling spend writes - verbose_proxy_logger.debug( - "_batch_database_updates: _record_request_error failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_org_transaction( - payload=payload_copy, - org_id=org_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_org_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_tag_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_tag_transaction failed: %s", - traceback.format_exc(), - ) - + except Exception as e: # noqa: BLE001 # the rollup must never skip the sibling spend writes + _report_spend_update_step_failure("_record_request_error", payload=payload_copy, exc=e) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_org_transaction", + self.add_spend_log_transaction_to_daily_org_transaction( + payload=payload_copy, org_id=org_id, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_tag_transaction", + self.add_spend_log_transaction_to_daily_tag_transaction(payload=payload_copy, prisma_client=prisma_client), + payload=payload_copy, + ) await self._enqueue_model_usage_transaction(payload=payload_copy, prisma_client=prisma_client) + @staticmethod + async def _run_spend_update_step( + step_name: str, step: Coroutine[object, object, object], *, payload: SpendLogsPayload + ) -> None: + try: + await step + except Exception as e: # noqa: BLE001 # one helper failing must not skip the sibling spend writes + _report_spend_update_step_failure(step_name, payload=payload, exc=e) + async def _update_key_db( self, response_cost: float | None, @@ -1657,15 +1626,13 @@ class DBSpendUpdateWriter: window_spend_update_queue=self.window_spend_update_queue, ) - await self._flush_tool_discovery_queue(prisma_client=prisma_client) - # Only commit from redis to db if this pod is the leader if await self.pod_lock_manager.acquire_lock( cronjob_id=DB_SPEND_UPDATE_JOB_NAME, ): verbose_proxy_logger.debug("acquired lock for spend updates") - uncommitted: dict[str, Any] = {} # mutable-ok: tracks popped categories still needing commit + uncommitted: dict[str, _SpendRestore] = {} # mutable-ok: tracks popped categories still needing commit committed_spend_tables: Final[list[_SpendTableName]] = [] # mutable-ok: filled as each table lands try: @@ -1680,13 +1647,41 @@ class DBSpendUpdateWriter: ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() uncommitted = { - "db_spend_update_transactions": db_spend_update_transactions, - "daily_spend_update_transactions": daily_spend_update_transactions, - "daily_team_spend_update_transactions": daily_team_spend_update_transactions, - "daily_org_spend_update_transactions": daily_org_spend_update_transactions, - "daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions, - "daily_agent_spend_update_transactions": daily_agent_spend_update_transactions, - "window_spend_update_transactions": window_spend_update_transactions, + "db_spend_update_transactions": ( + {} + if db_spend_update_transactions is None + else {"db_spend_update_transactions": db_spend_update_transactions} + ), + "daily_spend_update_transactions": ( + {} + if daily_spend_update_transactions is None + else {"daily_spend_update_transactions": daily_spend_update_transactions} + ), + "daily_team_spend_update_transactions": ( + {} + if daily_team_spend_update_transactions is None + else {"daily_team_spend_update_transactions": daily_team_spend_update_transactions} + ), + "daily_org_spend_update_transactions": ( + {} + if daily_org_spend_update_transactions is None + else {"daily_org_spend_update_transactions": daily_org_spend_update_transactions} + ), + "daily_end_user_spend_update_transactions": ( + {} + if daily_end_user_spend_update_transactions is None + else {"daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions} + ), + "daily_agent_spend_update_transactions": ( + {} + if daily_agent_spend_update_transactions is None + else {"daily_agent_spend_update_transactions": daily_agent_spend_update_transactions} + ), + "window_spend_update_transactions": ( + {} + if window_spend_update_transactions is None + else {"window_spend_update_transactions": window_spend_update_transactions} + ), } if db_spend_update_transactions is not None: @@ -1715,9 +1710,12 @@ class DBSpendUpdateWriter: on_table_committed=committed_spend_tables.append, ) except Exception as e: - uncommitted["db_spend_update_transactions"] = _spend_tables_left_to_send( + left_to_send: Final = _spend_tables_left_to_send( db_spend_update_transactions, committed_spend_tables, e ) + uncommitted["db_spend_update_transactions"] = ( + {} if left_to_send is None else {"db_spend_update_transactions": left_to_send} + ) raise uncommitted.pop("db_spend_update_transactions", None) @@ -1772,7 +1770,7 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: - if not _spend_commit_failure_is_requeue_safe(e): + if not db_write_failure_is_requeue_safe(e): uncommitted.pop("window_spend_update_transactions", None) spend_log_error( "Spend tracking - dropped %d budget window increments: the failed statement may have " @@ -1791,7 +1789,7 @@ class DBSpendUpdateWriter: exc=e, ) finally: - to_restore = {name: txns for name, txns in uncommitted.items() if txns is not None} + to_restore: Final = reduce(_merged_spend_restore, uncommitted.values(), _SpendRestore()) if to_restore: await self.redis_update_buffer.restore_transactions_to_redis(**to_restore) await self.pod_lock_manager.release_lock( @@ -1938,7 +1936,7 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run - if _spend_commit_failure_is_requeue_safe(e): + if db_write_failure_is_requeue_safe(e): spend_log_error( "Spend tracking - failed to commit budget window spend updates. " "Re-queued %d window increments for retry on next tick. Error: %s", @@ -1956,9 +1954,6 @@ class DBSpendUpdateWriter: exc=e, ) - ################## Tool Registry Upserts ################## - await self._flush_tool_discovery_queue(prisma_client=prisma_client) - async def commit_daily_tag_spend_to_db( self, prisma_client: PrismaClient, @@ -2090,20 +2085,6 @@ class DBSpendUpdateWriter: ) raise - async def _flush_tool_discovery_queue( - self, - prisma_client: PrismaClient, - ) -> None: - """Flush ToolDiscoveryQueue and batch-upsert new tools into LiteLLM_ToolTable.""" - from litellm.proxy.db.tool_registry_writer import batch_upsert_tools - - try: - items: Final = self.tool_discovery_queue.flush() - if items: - await batch_upsert_tools(prisma_client=prisma_client, items=items) - except Exception as e: - verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e) - @staticmethod async def handle_spend_update_failure( e: Exception, @@ -2115,7 +2096,6 @@ class DBSpendUpdateWriter: """Retry a failed spend-update transaction on connection errors, deadlocks or a rollup ``lock_timeout`` (55P03), else re-raise. All three roll the transaction back before any increment applied, so re-sending the same batch cannot double-count.""" - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.utils import raise_failed_update_spend_exception is_retryable = ( @@ -2624,7 +2604,7 @@ class DBSpendUpdateWriter: _mark_daily_spend_commit_started() _mark_daily_spend_commit_finished() except Exception as batch_error: - if _spend_commit_failure_is_requeue_safe(batch_error): + if db_write_failure_is_requeue_safe(batch_error): spend_log_error( "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", entity_type, @@ -2810,7 +2790,7 @@ class DBSpendUpdateWriter: date: Final = start_time.isoformat() if isinstance(start_time, datetime) else str(start_time or "") if not date: return - metadata: Final[SpendLogsMetadata] = json.loads(payload["metadata"]) + metadata: Final = _parse_spend_log_metadata(payload) internal_origin: Final = bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults request_error_accumulator.record( payload=payload, @@ -2861,7 +2841,7 @@ class DBSpendUpdateWriter: # spend, token and per-entity columns written here stay either way. request_status: Final = prisma_client.get_request_status(payload) verbose_proxy_logger.debug("Logged request status: %s", request_status) - _metadata: Final[SpendLogsMetadata] = json.loads(payload["metadata"]) + _metadata: Final = _parse_spend_log_metadata(payload) usage_obj: Final = _metadata.get("usage_object", {}) or {} if isinstance(payload["startTime"], datetime): start_time: Final = payload["startTime"].isoformat() diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 627b30e83f7..1116cd604a3 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -1,7 +1,7 @@ import asyncio import json import logging -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeAlias from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -14,7 +14,9 @@ from litellm.types.services import ServiceTypes POD_LOCK_TARGET: Final = "pod_lock" if TYPE_CHECKING: - ProxyLogging = Any + from litellm.proxy.utils import ProxyLogging as _ProxyLogging + + ProxyLogging: TypeAlias = _ProxyLogging else: ProxyLogging = Any diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 54d20bd4507..3a4375ac54a 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -9,8 +9,10 @@ import json from collections.abc import Mapping, Sequence from datetime import datetime from functools import reduce +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast +from pydantic import TypeAdapter from redis.exceptions import RedisError from litellm._internal_context import with_service_target @@ -26,6 +28,7 @@ from litellm.constants import ( REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, + REDIS_TOOL_DISCOVERY_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) @@ -41,12 +44,17 @@ from litellm.proxy._types import ( DBSpendUpdateTransactions, Litellm_EntityType, SpendUpdateQueueItem, + ToolDiscoveryQueueItem, ) from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, + aggregate_tool_discoveries, +) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, @@ -98,6 +106,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( ) _ValueT = TypeVar("_ValueT") +_DECODED_JSON: Final = TypeAdapter(object) def _spend_log_json_default(value: object) -> str: @@ -109,7 +118,7 @@ def _encode_spend_log_row(row: SpendLogRow) -> str: def _decode_spend_log_row(encoded: str) -> dict[str, object] | None: - decoded: Final = json.loads(encoded) + decoded: Final = _DECODED_JSON.validate_python(json.loads(encoded)) return decoded if isinstance(decoded, dict) else None @@ -133,6 +142,9 @@ def _merged_entity_transactions( ) +_TOOL_DISCOVERY_ITEMS: Final = TypeAdapter(list[ToolDiscoveryQueueItem]) + + class RedisUpdateBuffer: """ Handles buffering database `UPDATE` transactions in Redis before committing them to the database @@ -557,6 +569,56 @@ class RedisUpdateBuffer: str(e), ) + @with_service_target(SPEND_QUEUE_TARGET) + async def store_tool_discoveries_in_redis(self, tool_discovery_queue: ToolDiscoveryQueue) -> None: + """Move this worker's pending tool rows to Redis so only the leader writes LiteLLM_ToolTable.""" + if self.redis_cache is None: + return + items: Final = tool_discovery_queue.flush() + if not items: + return + try: + _ = await self.redis_cache.async_rpush_pipeline( + rpush_list=[ + RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(items)]) + ] + ) + except (RedisError, OSError, asyncio.TimeoutError) as e: + verbose_proxy_logger.error( + "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s", + len(items), + str(e), + ) + for item in items: + tool_discovery_queue.add_update(item) + + @with_service_target(SPEND_QUEUE_TARGET) + async def restore_tool_discoveries_to_redis(self, items: Sequence[ToolDiscoveryQueueItem]) -> None: + """Put drained tool rows back for the next drain after a write that provably did not apply.""" + if self.redis_cache is None or not items: + return + _ = await self.redis_cache.async_rpush_pipeline( + rpush_list=[ + RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(list(items))]) + ] + ) + + @with_service_target(SPEND_QUEUE_TARGET) + async def get_all_tool_discoveries_from_redis_buffer(self) -> tuple[ToolDiscoveryQueueItem, ...]: + """Drain every worker's tool rows and fold them to one row per tool_name.""" + if self.redis_cache is None: + return () + popped: Final = await self.redis_cache.async_lpop_pipeline( + lpop_list=[ + RedisPipelineLpopOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT) + ] + ) + entries: Final = popped[0] if popped else None + if not entries: + return () + batches: Final = tuple(_TOOL_DISCOVERY_ITEMS.validate_json(entry) for entry in entries) + return aggregate_tool_discoveries(chain.from_iterable(batches)) + @with_service_target(SPEND_QUEUE_TARGET) async def store_spend_logs_in_redis( self, diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index d5f6a447131..e9ec040023c 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -1,52 +1,57 @@ """ -In-memory buffer for tool registry upserts. - -Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue -uses set-deduplication: each unique tool_name is only queued once per flush -cycle (~30s). The seen-set is cleared on every flush so that call_count -increments in subsequent cycles rather than stopping after the first flush. +In-memory buffer for tool registry upserts, aggregated by tool_name so one flush carries one row per +tool with the number of times it was invoked since the previous flush. """ +from collections.abc import Iterable from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem +_ITEM: Final = TypeAdapter(ToolDiscoveryQueueItem) + + +def _merged(existing: ToolDiscoveryQueueItem | None, item: ToolDiscoveryQueueItem) -> ToolDiscoveryQueueItem: + """The first sighting keeps its origin and caller fields, later sightings only add to call_count.""" + increment: Final = item.get("call_count", 1) + if existing is None: + return _ITEM.validate_python({**item, "call_count": increment}) + return _ITEM.validate_python({**existing, "call_count": existing.get("call_count", 1) + increment}) + + +def aggregate_tool_discoveries(items: Iterable[ToolDiscoveryQueueItem]) -> tuple[ToolDiscoveryQueueItem, ...]: + """Collapse items to one per tool_name, summing call_count, in first-seen order.""" + by_name: Final[dict[str, ToolDiscoveryQueueItem]] = {} # mutable-ok: single pass fold keyed by tool_name + for item in items: + if not (tool_name := item.get("tool_name", "")): + continue + by_name[tool_name] = _merged(by_name.get(tool_name), item) + return tuple(by_name.values()) + class ToolDiscoveryQueue: - """ - In-memory buffer for tool registry upserts. - - Deduplicates by tool_name within each flush cycle: a tool is only queued - once per ~30s batch, so call_count increments once per flush cycle the - tool appears in (not once per invocation, but not once per pod lifetime - either). The seen-set is cleared on flush so subsequent batches can - re-count the same tool. - """ + """In-memory buffer for tool registry upserts, one pending row per tool_name with a running call_count.""" def __init__(self) -> None: - self._seen_tool_names: set[str] = set() - self._pending: list[ToolDiscoveryQueueItem] = [] + self._pending: dict[str, ToolDiscoveryQueueItem] = {} # mutable-ok: folded per add_update, cleared on flush def add_update(self, item: ToolDiscoveryQueueItem) -> None: - """Enqueue a tool discovery item if tool_name has not been seen before.""" + """Count one more invocation of item's tool_name, keeping the first sighting's origin and caller fields.""" tool_name: Final = item.get("tool_name", "") if not tool_name: return - if tool_name in self._seen_tool_names: - verbose_proxy_logger.debug("ToolDiscoveryQueue: skipping already-seen tool %s", tool_name) - return - self._seen_tool_names.add(tool_name) - self._pending.append(item) + self._pending[tool_name] = _merged(self._pending.get(tool_name), item) verbose_proxy_logger.debug( - "ToolDiscoveryQueue: queued new tool %s (origin=%s)", + "ToolDiscoveryQueue: tool %s call_count=%s (origin=%s)", tool_name, + self._pending[tool_name].get("call_count"), item.get("origin"), ) def flush(self) -> list[ToolDiscoveryQueueItem]: - """Return and clear all pending items. Resets seen-set so the next - flush cycle can re-count the same tools.""" - items, self._pending = self._pending, [] - self._seen_tool_names.clear() + """Return and clear all pending rows.""" + items, self._pending = list(self._pending.values()), {} return items diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 379e0012ea8..5dfcaa4b2f8 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -8,6 +8,7 @@ from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, ProxyErrorTypes, ProxyException, ) @@ -637,3 +638,14 @@ async def call_with_db_reconnect_retry( # propagate — repeated reconnect-loops are the watchdog's job, not # this helper's. return await coro_factory() + + +_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) + + +def db_write_failure_is_requeue_safe(e: Exception) -> bool: + """Whether a failed write provably never applied, so its rows can be sent again without double counting.""" + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES diff --git a/litellm/proxy/db/prisma_query_span.py b/litellm/proxy/db/prisma_query_span.py index 75625c6c0f2..2747a765f2a 100644 --- a/litellm/proxy/db/prisma_query_span.py +++ b/litellm/proxy/db/prisma_query_span.py @@ -20,6 +20,8 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + from litellm.integrations.otel.model.spans import PG_CATALOG, PRISMA_RELATIONS @@ -34,6 +36,7 @@ class PrismaQuery: UNKNOWN_PRISMA_QUERY: Final = PrismaQuery("prisma_query", None, None) +_DECODED_JSON: Final = TypeAdapter(object) _ROOT_FIELD: Final = re.compile(r"result:\s*(\w+)") _RAW_SQL: Final = re.compile(r'query:\s*"((?:[^"\\]|\\.)*)"') @@ -111,7 +114,7 @@ def sql_operation(sql: str) -> tuple[str | None, str | None]: def _query_text(content: str) -> str: try: - payload: Final[object] = json.loads(content) + payload: Final = _DECODED_JSON.validate_python(json.loads(content)) except ValueError: return content query: Final = payload.get("query") if isinstance(payload, dict) else None diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py index 673b549b645..3214c0974b4 100644 --- a/litellm/proxy/db/proxy_worker_heartbeat.py +++ b/litellm/proxy/db/proxy_worker_heartbeat.py @@ -56,6 +56,7 @@ class _LiveWorkerCountRow(TypedDict): _COUNT_ROWS_ADAPTER: Final = TypeAdapter(tuple[_LiveWorkerCountRow, ...]) +_RAW_QUERY_RESULT: Final = TypeAdapter(object) class _HeartbeatExecute(Protocol): @@ -99,7 +100,9 @@ async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None: db: Final = prisma_client.db primary_db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) else db async with db_span("count_live_proxy_workers", "LiteLLM_ProxyWorkerHeartbeat"): - rows: Final = await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + rows: Final = _RAW_QUERY_RESULT.validate_python( + await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + ) return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"] except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503 verbose_proxy_logger.debug("Live proxy worker count unavailable: %s", count_err) diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 48ad57fd06e..52a0be3d49e 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -298,7 +298,7 @@ class RoutingPrismaWrapper: def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: - return getattr(self.read_target, name) + return _DELEGATED_ATTRIBUTE.validate_python(getattr(self.read_target, name)) writer_attr: Final = _DELEGATED_ATTRIBUTE.validate_python(getattr(self._writer, name)) # Per-model action accessors are non-callable instances that expose # both `find_many` and `create`. Methods like execute_raw / batch_ / diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 7ac414bec00..591ad81a4e5 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -19,7 +19,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from itertools import chain, groupby -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from typing import TYPE_CHECKING, Final, Literal, Protocol from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import LiteralString @@ -122,7 +122,7 @@ def build_tool_usage_transaction( spend: float, total_tokens: int, completion_response: object, - realtime_tool_calls: Any = None, + realtime_tool_calls: object = None, ) -> ToolUsageTransaction | None: """None when the request invoked no tools. Realtime sessions carry invoked tools in kwargs["realtime_tool_calls"] (OpenAI tool_calls shape) rather than diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 2e3530d7613..f7ec88945a1 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -14,8 +14,17 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.constants import TOOL_DISCOVERY_COMMIT_JOB_NAME from litellm.proxy._types import ToolDiscoveryQueueItem -from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, + aggregate_tool_discoveries, +) +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe +from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ToolRepository @@ -105,58 +114,75 @@ def _row_to_model(row: object) -> LiteLLM_ToolTableRow: async def batch_upsert_tools( prisma_client: "PrismaClient", - items: list[ToolDiscoveryQueueItem], + items: Sequence[ToolDiscoveryQueueItem], ) -> None: - """ - Batch-upsert tool registry rows via Prisma. + """Write a batch of discovered tools to LiteLLM_ToolTable in a single statement. Raises on failure so the + caller decides whether the rows can be retried.""" + rows: Final = aggregate_tool_discoveries(items) + if not rows: + return + sql, params = build_tool_bulk_upsert(rows) + async with db_span("commit_tool_discoveries", "LiteLLM_ToolTable"): + await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client + verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(rows)) - On first insert: sets input_policy/output_policy = "untrusted" (default), call_count = 1. - On conflict: increments call_count; preserves existing policies. - """ + +async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discovery_queue: ToolDiscoveryQueue) -> None: + """One worker's own rows straight to LiteLLM_ToolTable, the path without the Redis buffer.""" + items: Final = tool_discovery_queue.flush() if not items: return try: - data: Final = [item for item in items if item.get("tool_name")] - if not data: - return - now: Final = datetime.now(timezone.utc) - table: Final = _tool_table_actions(prisma_client) - for item in data: - tool_name = item.get("tool_name", "") - origin = item.get("origin") or "user_defined" - created_by = item.get("created_by") or "system" - key_hash = item.get("key_hash") - team_id = item.get("team_id") - key_alias = item.get("key_alias") - user_agent = item.get("user_agent") - await table.upsert( - where={"tool_name": tool_name}, - data={ - "create": { - "tool_id": str(uuid.uuid4()), - "tool_name": tool_name, - "origin": origin, - "input_policy": "untrusted", - "output_policy": "untrusted", - "call_count": 1, - "created_by": created_by, - "updated_by": created_by, - "key_hash": key_hash, - "team_id": team_id, - "key_alias": key_alias, - "user_agent": user_agent, - "last_used_at": now, - }, - "update": { - "call_count": {"increment": 1}, - "updated_at": now, - "last_used_at": now, - }, - }, + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail + verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e) + + +async def commit_tool_discoveries( + prisma_client: "PrismaClient", + tool_discovery_queue: ToolDiscoveryQueue, + redis_update_buffer: RedisUpdateBuffer, + pod_lock_manager: PodLockManager, + use_redis_buffer: bool, +) -> None: + """Scheduled alongside update_spend on its own lock, so tool table writes never ride the spend commit. + + With the Redis buffer every worker hands its rows to Redis and the holder of the tool lock writes them all in + one statement. Without it each worker writes its own rows.""" + if not use_redis_buffer: + await flush_tool_discovery_queue(prisma_client=prisma_client, tool_discovery_queue=tool_discovery_queue) + return + await redis_update_buffer.store_tool_discoveries_in_redis(tool_discovery_queue=tool_discovery_queue) + if not await pod_lock_manager.acquire_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME): + return + try: + await _commit_tool_discoveries_from_redis(prisma_client=prisma_client, redis_update_buffer=redis_update_buffer) + finally: + await pod_lock_manager.release_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME) + + +async def _commit_tool_discoveries_from_redis( + prisma_client: "PrismaClient", redis_update_buffer: RedisUpdateBuffer +) -> None: + """Tool lock holder only: drain every worker's tool rows from Redis and write them in one statement.""" + items: Final = await redis_update_buffer.get_all_tool_discoveries_from_redis_buffer() + if not items: + return + try: + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: # noqa: BLE001 # the rows go back to redis or are dropped, the scheduler job must not fail + if db_write_failure_is_requeue_safe(e): + await redis_update_buffer.restore_tool_discoveries_to_redis(items) + verbose_proxy_logger.error( + "Tool discovery - %d tool rows went back to Redis for the next tick. Error: %s", len(items), str(e) ) - verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(data)) - except Exception as e: - verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e) + return + verbose_proxy_logger.error( + "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database " + "refused the data, so re-sending it is not safe. Error: %s", + len(items), + str(e), + ) _NO_OWNERS: Final[Mapping[str, ToolDiscoveryUser]] = MappingProxyType({}) diff --git a/litellm/proxy/db/tool_table_bulk_upsert.py b/litellm/proxy/db/tool_table_bulk_upsert.py new file mode 100644 index 00000000000..e586042ae7c --- /dev/null +++ b/litellm/proxy/db/tool_table_bulk_upsert.py @@ -0,0 +1,71 @@ +"""One multi-row ``INSERT ... ON CONFLICT ("tool_name") DO UPDATE`` per batch of discovered tools. + +LiteLLM_ToolTable holds one row per tool name the proxy has seen. A batch arrives already folded to one +entry per tool name with a summed call_count, and this statement inserts the names the table has not seen +with the untrusted policies the columns default to, and for the names it has adds the batch's call_count +while leaving the policies and the discovering caller's fields alone. Same shape as daily_spend_bulk_upsert. +""" + +import uuid +from collections.abc import Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.proxy._types import ToolDiscoveryQueueItem + +SqlValue: TypeAlias = str | int | None + +_TOOL_ROW_COLUMNS: Final[tuple[str, ...]] = ( + "tool_id", + "tool_name", + "origin", + "call_count", + "created_by", + "updated_by", + "key_hash", + "team_id", + "key_alias", + "user_agent", +) +_CASTS: Final[Mapping[str, str]] = MappingProxyType({"call_count": "integer"}) + + +def _row_params(item: ToolDiscoveryQueueItem) -> tuple[SqlValue, ...]: + created_by: Final = item.get("created_by") or "system" + return ( + str(uuid.uuid4()), + item.get("tool_name", ""), + item.get("origin") or "user_defined", + item.get("call_count", 1), + created_by, + created_by, + item.get("key_hash"), + item.get("team_id"), + item.get("key_alias"), + item.get("user_agent"), + ) + + +def build_tool_bulk_upsert(items: Sequence[ToolDiscoveryQueueItem]) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one batch of distinct tool_names, plus its positional arguments.""" + width: Final = len(_TOOL_ROW_COLUMNS) + rows: Final = ", ".join( + "(" + + ", ".join( + f"${row_index * width + offset + 1}::{_CASTS.get(column, 'text')}" + for offset, column in enumerate(_TOOL_ROW_COLUMNS) + ) + + ", (NOW() AT TIME ZONE 'UTC'), (NOW() AT TIME ZONE 'UTC'))" + for row_index in range(len(items)) + ) + columns: Final = ", ".join(f'"{column}"' for column in _TOOL_ROW_COLUMNS) + sql: Final = ( + f'INSERT INTO "LiteLLM_ToolTable" ({columns}, "last_used_at", "updated_at")\n' + f"VALUES {rows}\n" + 'ON CONFLICT ("tool_name") DO UPDATE SET\n' + ' "call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count",\n' + ' "last_used_at" = EXCLUDED."last_used_at",\n' + ' "updated_at" = EXCLUDED."updated_at"' + ) + return sql, tuple(chain.from_iterable(map(_row_params, items))) diff --git a/litellm/proxy/decisions_endpoints/endpoints.py b/litellm/proxy/decisions_endpoints/endpoints.py index 366842d23a6..a8684dd9757 100644 --- a/litellm/proxy/decisions_endpoints/endpoints.py +++ b/litellm/proxy/decisions_endpoints/endpoints.py @@ -228,7 +228,7 @@ async def systemone( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ - Judge a `state` against 1 to 128 named `questions` and get one calibrated answer per question name. + Judge a `state` against one or more named `questions` and get one calibrated answer per question name. Question types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with `confidence` and `probabilities`), `score` (an index into the `criteria` levels, with `confidence`, diff --git a/litellm/proxy/discovery_endpoints/agent_skills_archive.py b/litellm/proxy/discovery_endpoints/agent_skills_archive.py index 1f2fca3992e..94a0c3e2cf1 100644 --- a/litellm/proxy/discovery_endpoints/agent_skills_archive.py +++ b/litellm/proxy/discovery_endpoints/agent_skills_archive.py @@ -16,6 +16,7 @@ from types import MappingProxyType from typing import Final import yaml +from pydantic import TypeAdapter MAX_ARCHIVE_UNPACKED_BYTES: Final = 50 * 1024 * 1024 MAX_ARCHIVE_ENTRIES: Final = 1000 @@ -26,6 +27,7 @@ _ZIP_ENTRY_PERMISSIONS: Final = 0o644 << 16 _FRONTMATTER_PATTERN: Final = re.compile(r"^---\s*\n(.*?)\n---\s*(?:\n|$)", re.DOTALL) _WINDOWS_DRIVE_PATTERN: Final = re.compile(r"^[A-Za-z]:") _EMPTY_FRONTMATTER: Final[Mapping[str, object]] = MappingProxyType({}) +_DECODED_YAML: Final = TypeAdapter(object) @dataclass(frozen=True, slots=True) @@ -100,7 +102,7 @@ def _manifest_frontmatter(manifest: bytes) -> Mapping[str, object]: if match is None: return _EMPTY_FRONTMATTER try: - parsed: Final = yaml.safe_load(match.group(1)) + parsed: Final = _DECODED_YAML.validate_python(yaml.safe_load(match.group(1))) except yaml.YAMLError: return _EMPTY_FRONTMATTER if not isinstance(parsed, dict): diff --git a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py index 2050cc40e6d..830f7b0334e 100644 --- a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py +++ b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py @@ -10,7 +10,6 @@ import asyncio import re from collections.abc import Sequence from itertools import groupby -from operator import itemgetter from types import MappingProxyType from typing import Final @@ -173,11 +172,15 @@ def _slugify(raw: str | None) -> str | None: return _NON_SLUG_PATTERN.sub("-", raw.lower()).strip("-")[:MAX_SKILL_NAME_LENGTH].rstrip("-") or None +def _entry_name(entry: tuple[int, str]) -> str: + return entry[1] + + def _deduplicated(names: Sequence[str]) -> tuple[str, ...]: ordinals: Final = MappingProxyType( { position: ordinal - for _, duplicates in groupby(sorted(enumerate(names), key=itemgetter(1)), key=itemgetter(1)) + for _, duplicates in groupby(sorted(enumerate(names), key=_entry_name), key=_entry_name) for ordinal, (position, _) in enumerate(duplicates) } ) diff --git a/litellm/proxy/example_config_yaml/custom_auth.py b/litellm/proxy/example_config_yaml/custom_auth.py index b7646ce5e3d..c53074638c6 100644 --- a/litellm/proxy/example_config_yaml/custom_auth.py +++ b/litellm/proxy/example_config_yaml/custom_auth.py @@ -30,7 +30,7 @@ async def generate_key_fn(data: GenerateKeyRequest): data_json: Final = data.json() # Unpacking variables - team_id: Final = data_json.get("team_id") + team_id: Final = data.team_id data_json.get("duration") data_json.get("models") data_json.get("aliases") diff --git a/litellm/proxy/example_config_yaml/custom_callbacks.py b/litellm/proxy/example_config_yaml/custom_callbacks.py index 2c689811dc8..5ae55060930 100644 --- a/litellm/proxy/example_config_yaml/custom_callbacks.py +++ b/litellm/proxy/example_config_yaml/custom_callbacks.py @@ -7,10 +7,14 @@ import traceback sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import inspect +from typing import Final + +from pydantic import TypeAdapter import litellm from litellm.integrations.custom_logger import CustomLogger -from typing import Final + +_ATTRIBUTE_VALUE: Final = TypeAdapter(object) # This file includes the custom callbacks for LiteLLM Proxy @@ -27,7 +31,11 @@ class MyCustomHandler(CustomLogger): print_verbose(f"{blue_color_code}Initialized LiteLLM custom logger") try: print_verbose("Logger Initialized with following methods:") - methods: Final = [method for method in dir(self) if inspect.ismethod(getattr(self, method))] + methods: Final = [ + method + for method in dir(self) + if inspect.ismethod(_ATTRIBUTE_VALUE.validate_python(getattr(self, method))) + ] # Pretty print_verbose the methods for method in methods: diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 64c935ce0d0..8581dca603b 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -9,7 +9,14 @@ every text fragment. """ from collections.abc import Callable, Iterator, Mapping, Sequence -from typing import Any, Final +from typing import Final + +from litellm.litellm_core_utils.core_helpers import iter_items +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + input_items, + with_visited_compaction_summary, +) # Call types whose body carries free-form chat / prompt text that # text-content guardrails (banned keywords, content moderation, secret @@ -95,12 +102,12 @@ def _coerce_input_to_messages(input_value: object) -> list[dict[str, object]]: """Coerce a Responses-API ``data["input"]`` value into chat-style messages.""" if isinstance(input_value, str): return [{"role": "user", "content": input_value}] - if not isinstance(input_value, list): - return [] messages: Final[list[dict[str, object]]] = [] - for item in input_value: + for item in input_items(input_value): if isinstance(item, str): messages.append({"role": "user", "content": item}) + elif (summary := compaction_summary(item)) is not None: + messages.append({"role": "user", "content": summary}) elif isinstance(item, dict): if _part_text(item) is not None: messages.append({"role": item.get("role") or "user", "content": [item]}) @@ -189,11 +196,14 @@ def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int: data["input"] = visit(input_value) return visited if isinstance(input_value, list): - for idx, item in enumerate(input_value): + for idx, item in enumerate(input_items(input_value)): if isinstance(item, str): if item: visited += 1 input_value[idx] = visit(item) + elif (compaction := with_visited_compaction_summary(item, visit)) is not None: + visited += 1 + input_value[idx] = compaction elif isinstance(item, dict): if _part_text(item) is not None: visited += 1 @@ -222,7 +232,7 @@ def is_string_batch_input(data: Mapping[str, object]) -> bool: return isinstance(input_value, list) and bool(input_value) and all(isinstance(item, str) for item in input_value) -def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: Sequence[object]) -> bool: +def apply_redacted_messages_back(data: dict[str, object], redacted_messages: Sequence[object]) -> bool: """Write redacted messages back to whichever field(s) the caller used. Mask/anonymize paths take a synthesised messages list (from @@ -240,9 +250,9 @@ def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: Sequen carry exactly one message per inspected element: a partial rewrite would forward the remaining originals unredacted. Callers must block on False. """ + input_value: Final = data.get("input") if is_string_batch_input(data): - batch: Final = data["input"] - inspected_indices: Final = tuple(idx for idx, item in enumerate(batch) if item) + inspected_indices: Final = tuple(idx for idx, item in enumerate(iter_items(input_value)) if item) if len(redacted_messages) != len(inspected_indices): return False if any(not isinstance(message, Mapping) or message.get("content") is None for message in redacted_messages): @@ -250,12 +260,12 @@ def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: Sequen redacted_texts: Final = tuple( "\n".join(_iter_text_parts_in_content(message["content"])) for message in redacted_messages ) - for idx, text in zip(inspected_indices, redacted_texts): - batch[idx] = text + if isinstance(input_value, list): + for idx, text in zip(inspected_indices, redacted_texts): + input_value[idx] = text return True if "messages" in data: data["messages"] = redacted_messages - input_value: Final = data.get("input") if isinstance(input_value, str): text_parts: Final[list[str]] = [] for msg in redacted_messages: diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 70e7aac12b5..a5822bfb0e3 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -90,6 +90,7 @@ if TYPE_CHECKING: router: Final = APIRouter() GUARDRAIL_REGISTRY: Final = GuardrailRegistry() _DECODED_JSON: Final = TypeAdapter(object) +_DECODED_YAML: Final = TypeAdapter(object) def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]: @@ -1682,7 +1683,7 @@ async def validate_blocked_words_file(request: dict[str, str]): if not file_content: return {"valid": False, "error": "No file content provided"} - data: Final = yaml.safe_load(file_content) + data: Final = _DECODED_YAML.validate_python(yaml.safe_load(file_content)) if not isinstance(data, dict) or "blocked_words" not in data: return { diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py index ac0a6911184..aabac874ca2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -63,6 +63,7 @@ _INVALID_ASSERTION_AADSTS_PREFIX: Final = "50027" _AADSTS_CODES_ADAPTER: Final = TypeAdapter(tuple[int, ...]) _MCP_CALL_TYPES: Final[tuple[str, ...]] = ("mcp_call", "call_mcp_tool") _TOOL_INPUT_SCHEMA_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) _OBO_CACHE_MAX_ENTRIES: Final = 1000 _DEFAULT_TOKEN_TTL_SECONDS: Final = 3599.0 _TOKEN_EXPIRY_SLACK_SECONDS: Final = 60.0 @@ -343,7 +344,7 @@ class Agent365Guardrail(CustomGuardrail): latency_ms: float, ) -> dict: # mutable-ok: returns the request data dict per hook contract try: - parsed_verdict: Final = response.json() + parsed_verdict: Final = _DECODED_JSON.validate_python(response.json()) except ValueError: return self._handle_unavailable( data=data, @@ -486,7 +487,7 @@ class Agent365Guardrail(CustomGuardrail): response=response, ) try: - parsed_body: Final = response.json() + parsed_body: Final = _DECODED_JSON.validate_python(response.json()) except ValueError as exc: raise Agent365MalformedResponseError("the Entra token endpoint returned a non-JSON body") from exc if not isinstance(parsed_body, dict): diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py index 993c1d30949..389706d7417 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -2,12 +2,13 @@ import os import uuid from collections.abc import Mapping from typing import ( - Any, # noqa: TID251 # the only type CustomGuardrail.__init__ accepts for its open-ended kwargs Final, Literal, ) import httpx +from pydantic import TypeAdapter +from typing_extensions import TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException @@ -36,6 +37,12 @@ DEFAULT_BLOCKED_MESSAGE: Final = "Blocked by your organization's content policy. SUPPORTED_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call) +_DECODED_JSON: Final = TypeAdapter(object) + + +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + pass + def _stream_cannot_carry_a_rewrite(request_data: Mapping[str, object] | None) -> bool: """Whether this is a streamed response on a route the framework can only block, not rewrite. @@ -97,7 +104,7 @@ class AiriaGuardrail(CustomGuardrail): timeout: float | None = None, supported_event_hooks: list[GuardrailEventHooks] | None = None, # mutable-ok: matches CustomGuardrail.__init__ async_handler: AsyncHTTPHandler | None = None, # dependency-injected in tests; defaults to a real client - **kwargs: Any, # kwargs-ok: passed straight through to CustomGuardrail.__init__ + **kwargs: Unpack[_CustomGuardrailOptions], # kwargs-ok: passed straight through to CustomGuardrail.__init__ ) -> None: resolved_timeout: Final = timeout or float(os.getenv("AIRIA_TIMEOUT", "10")) @@ -161,7 +168,7 @@ class AiriaGuardrail(CustomGuardrail): headers={"Authorization": f"Bearer {self.api_key}"}, ) response.raise_for_status() - raw_body: Final[object] = response.json() + raw_body: Final = _DECODED_JSON.validate_python(response.json()) body: Final = _as_object(raw_body) if body is None: raise TypeError(f"expected a JSON object in the response body, got {type(raw_body).__name__}") diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py index 64770fdf0f7..5bf471265c8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations from .block_code_execution import BlockCodeExecutionGuardrail @@ -26,7 +27,7 @@ def _get_param( value: Final = getattr(litellm_params, key, default) if value is not None: return value - raw: Final = guardrail.get("litellm_params") + raw: Final = as_str_mapping(guardrail.get("litellm_params")) if isinstance(raw, dict) and key in raw: return raw[key] return default diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 936f862b10b..19b9bccd8aa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -23,6 +23,7 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_ssl_configuration, @@ -222,7 +223,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return [] @staticmethod - def _iter_schema_string_refs(data: Mapping[str, Any]): + def _iter_schema_string_refs(data: Mapping[str, object]): """Yield ``(container, key)`` for every non-empty schema string the proxy forwards to the model inside tool/function and structured-output schemas: each ``tools[].function`` and legacy ``functions[]`` entry plus the @@ -235,10 +236,10 @@ class CatoNetworksGuardrail(CustomGuardrail): list_keys: Final = ("enum", "examples") stack: Final[list] = [] - for tool in data.get("tools") or []: + for tool in iter_items(data.get("tools") or ()): if isinstance(tool, dict) and isinstance(tool.get("function"), dict): stack.append(tool["function"]) - for function in data.get("functions") or []: + for function in iter_items(data.get("functions") or ()): if isinstance(function, dict): stack.append(function) response_format: Final = data.get("response_format") diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index ab43d32e260..45e785fd53c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -35,6 +35,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -1743,7 +1744,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_masked_entity_count( - rules: list[dict[str, Any]], + rules: Sequence[Mapping[str, object]], ) -> dict[str, int] | None: """Count entity-type detections per Cisco rule for the logging payload.""" if not rules: @@ -1752,8 +1753,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): for rule in rules: if not isinstance(rule, dict): continue - entity_types = rule.get("entity_types") or [] - for entity_type in entity_types: + for entity_type in iter_items(rule.get("entity_types") or ()): if not isinstance(entity_type, str): continue counts[entity_type] = counts.get(entity_type, 0) + 1 diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py index 498f9bf4099..c5ca97ff760 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Final +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.types.guardrails import ( GuardrailEventHooks, Mode, @@ -26,7 +27,7 @@ def _coerce_event_hook( def _get_optional_value(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> object: if optional_params is not None: - value: Final = getattr(optional_params, attribute_name, None) + value: Final = attribute_or_none(optional_params, attribute_name) if value is not None: return value return getattr(litellm_params, attribute_name, None) diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 029feba89fb..4d910cde603 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -28,6 +28,7 @@ from urllib.parse import urlparse import httpx from fastapi import HTTPException from httpx import Response as HttpxResponse +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -72,6 +73,7 @@ if TYPE_CHECKING: ) BYPASS_HEADER: Final = "x-compresr-bypass" +_AGENTIC_LOOP_MODEL: Final = TypeAdapter(object) COMPRESR_RETRIEVE_TOOL_NAME: Final = "compresr_retrieve" DEFAULT_API_BASE: Final = "https://api.compresr.ai" DEFAULT_COMPRESSION_MODEL: Final = "latte_v2" @@ -1118,7 +1120,7 @@ class CompresrGuardrail(CustomGuardrail): full_model_name = model if logging_obj is not None: agentic_params: Final = getattr(logging_obj, "model_call_details", {}).get("agentic_loop_params", {}) - candidate: Final = agentic_params.get("model", model) + candidate: Final = _AGENTIC_LOOP_MODEL.validate_python(agentic_params.get("model", model)) if isinstance(candidate, str) and candidate: full_model_name = candidate diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 34725803c60..ee3a957f607 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -5,7 +5,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple, Optional, cast from fastapi import HTTPException -from pydantic import ConfigDict, Field, ValidationError +from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import override from litellm._logging import verbose_proxy_logger @@ -37,6 +37,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + class CrowdStrikeAIDRGuardrailMissingSecrets(Exception): """Custom exception for missing CrowdStrike AIDR secrets.""" @@ -362,7 +364,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): assert response is not None response.raise_for_status() - response_body: Final[object] = response.json() + response_body: Final = _DECODED_JSON.validate_python(response.json()) raw_result: Final[object] = response_body.get("result") if isinstance(response_body, dict) else None blocked_signal: Final[object] = raw_result.get("blocked") if isinstance(raw_result, dict) else None diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py index 13d7178dfe9..0f9ce15e65b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -39,7 +39,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" raise ValueError("Custom code guardrail requires a guardrail_name") # Get the custom code from litellm_params - custom_code: Final = getattr(litellm_params, "custom_code", None) + custom_code: Final = litellm_params.custom_code if not custom_code: raise ValueError("Custom code guardrail requires 'custom_code' in litellm_params") diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 199da7df30d..ecd23d86540 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, @@ -126,6 +127,13 @@ def _sanitize_inbound_headers( return sanitized or None +def _logged_inbound_headers(details: Mapping[str, object]) -> Mapping[object, object] | None: + litellm_params: Final = as_str_mapping(details.get("litellm_params", {})) + metadata: Final = as_str_mapping(litellm_params.get("metadata", {})) if litellm_params is not None else None + headers: Final = metadata.get("headers", None) if metadata is not None else None + return headers if isinstance(headers, dict) else None + + def _extract_inbound_headers( request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"], @@ -163,7 +171,7 @@ def _extract_inbound_headers( if logging_obj and getattr(logging_obj, "model_call_details", None): try: details: Final = logging_obj.model_call_details or {} - headers = details.get("litellm_params", {}).get("metadata", {}).get("headers", None) + headers = _logged_inbound_headers(details) if headers: return _sanitize_inbound_headers(headers, extra_allowlist=extra_allowlist) except Exception: @@ -283,7 +291,7 @@ class GenericGuardrailAPI(CustomGuardrail): ) self.num_retries: int = num_retries or 0 self.headers = headers or {} - self.extra_headers = extra_headers or [] + self.extra_headers: Sequence[object] = extra_headers or [] # If api_key is provided, add it as x-api-key header if api_key: diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index f32228a6204..7ff5aa2466b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj GRAYSWAN_BLOCK_ERROR_MSG: Final = "Blocked by Gray Swan Guardrail" +_DECODED_JSON: Final = TypeAdapter(object) class _GraySwanMonitorResponse(TypedDict): @@ -705,7 +706,9 @@ class GraySwanGuardrail(CustomGuardrail): {**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers ) if cleaned_litellm_metadata: - sanitized: Final[object] = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + sanitized: Final = _DECODED_JSON.validate_python( + safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + ) if isinstance(sanitized, dict) and sanitized: payload["litellm_metadata"] = sanitized diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py index d5edcc19a02..a1d0900ebed 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py +++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py @@ -2,6 +2,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Final from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -23,6 +24,8 @@ from litellm.types.utils import CallTypesLiteral, GuardrailStatus if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + class JavelinGuardrail(CustomGuardrail): @classmethod @@ -113,7 +116,9 @@ class JavelinGuardrail(CustomGuardrail): json=dict(request), timeout=self.timeout, ) - verbose_proxy_logger.debug("Javelin Guardrail: Javelin guard API response: %s", response.json()) + verbose_proxy_logger.debug( + "Javelin Guardrail: Javelin guard API response: %s", _DECODED_JSON.validate_python(response.json()) + ) response_data: Final = response.json() # Ensure the response has the required assessments field if "assessments" not in response_data: diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 8efd1dc79e3..7c0cdb0d36f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -7,6 +7,7 @@ from types import MappingProxyType from typing import Final, Literal from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -40,6 +41,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( ) from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponse +_DECODED_JSON: Final = TypeAdapter(object) +_RESPONSE_FIELDS: Final = TypeAdapter(Mapping[str, object]) + _DETECTOR_CATEGORY_PHRASES: Final[Mapping[str, str]] = MappingProxyType( { "prompt_injection": "a potential prompt injection attempt", @@ -402,7 +406,9 @@ class LakeraAIGuardrail(CustomGuardrail): json=request, timeout=self.timeout, ) - verbose_proxy_logger.debug("Lakera AI v2 guard response: %s", response.json()) + verbose_proxy_logger.debug( + "Lakera AI v2 guard response: %s", _DECODED_JSON.validate_python(response.json()) + ) lakera_response = LakeraAIResponse(**response.json()) return lakera_response, masked_entity_count except Exception as e: @@ -801,7 +807,7 @@ class LakeraAIGuardrail(CustomGuardrail): choice_idx = choice_indices[idx] response_dict["choices"][choice_idx]["message"]["content"] = msg.get("content", "") add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) - return ModelResponse(**response_dict) + return ModelResponse(**_RESPONSE_FIELDS.validate_python(response_dict)) # inject_system_message has nothing left to inject into once a response # already exists, so it is treated the same as monitor: log and allow. diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index e05db5003fa..97b35769f63 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -9,7 +9,7 @@ import json import os import uuid from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict +from typing import TYPE_CHECKING, Final, Literal, TypedDict try: import ulid @@ -39,6 +39,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.integrations.custom_guardrail import dc as global_cache +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -397,13 +398,15 @@ class LassoGuardrail(CustomGuardrail): HTTPException: If blocking violations are detected """ raw_messages: Final[list[dict[str, object]]] = data.get("messages") or [] - messages: list[dict[str, Any]] = self._expand_messages_for_classification(raw_messages) if raw_messages else [] - messages_count: Final = len(messages) - if data.get("input") is not None: - # Responses-API payloads carry text in data["input"]. Inspect it - # alongside any "messages" array — otherwise a caller can attach - # benign messages and stash blocked content in input to bypass. - messages.extend(build_inspection_messages({"input": data["input"]})) + expanded_messages: Final = self._expand_messages_for_classification(raw_messages) if raw_messages else [] + messages_count: Final = len(expanded_messages) + # Responses-API payloads carry text in data["input"]. Inspect it + # alongside any "messages" array — otherwise a caller can attach + # benign messages and stash blocked content in input to bypass. + input_messages: Final = ( + build_inspection_messages({"input": data["input"]}) if data.get("input") is not None else [] + ) + messages: Final[Sequence[Mapping[str, object]]] = [*expanded_messages, *input_messages] if not messages: return data @@ -421,7 +424,7 @@ class LassoGuardrail(CustomGuardrail): data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"], - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], ) -> dict: """Handle classification without masking.""" try: @@ -439,7 +442,7 @@ class LassoGuardrail(CustomGuardrail): data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"], - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], messages_count: int, ) -> dict: """Handle masking with classifix endpoint. @@ -470,7 +473,7 @@ class LassoGuardrail(CustomGuardrail): text_parts = [msg["content"] for msg in masked_for_input if isinstance(msg.get("content"), str)] if text_parts: data["input"] = "\n".join(text_parts) - self._log_masking_applied(message_type, dict(response)) + self._log_masking_applied(message_type, response) return data except Exception as e: @@ -479,7 +482,7 @@ class LassoGuardrail(CustomGuardrail): def _map_masked_messages_back( self, - original_messages: list[dict[str, Any]], + original_messages: Sequence[Mapping[str, object]], masked_messages: Sequence[Mapping[str, object]], ) -> list[dict[str, object]]: """Map Lasso-format masked messages back onto the original OpenAI-format messages. @@ -492,7 +495,7 @@ class LassoGuardrail(CustomGuardrail): """ # Index masked content by type so we can look up by id without caring about order. masked_tool_use: Final[dict[object, dict[str, object]]] = {} - masked_tool_result: Final[dict[str, str]] = {} + masked_tool_result: Final[dict[object, str]] = {} masked_text: Final[list[str]] = [] for msg in masked_messages: @@ -567,12 +570,12 @@ class LassoGuardrail(CustomGuardrail): def _update_tool_calls_from_masked( self, - tool_calls: list[object], + tool_calls: object, masked_tool_use: Mapping[object, Mapping[str, object]], ) -> list[object]: """Replace tool_call arguments with masked values returned by Lasso.""" updated: Final = [] - for call in tool_calls: + for call in iter_items(tool_calls): call_id = self._get_field(call, "id") if call_id and call_id in masked_tool_use: masked_input = masked_tool_use[call_id].get("input") @@ -627,7 +630,7 @@ class LassoGuardrail(CustomGuardrail): def _log_masking_applied( self, message_type: Literal["PROMPT", "COMPLETION"], - response: dict[str, Any], + response: LassoResponse, ) -> None: """Log masking application with structured context.""" conversation_id: Final = getattr(self, "conversation_id", "unknown") @@ -637,12 +640,12 @@ class LassoGuardrail(CustomGuardrail): "guardrail_name": getattr(self, "guardrail_name", "unknown"), "message_type": message_type, "violations_count": len(response.get("findings", {})), - "masked_fields": len(response.get("messages", [])), + "masked_fields": len(response.get("messages") or []), "conversation_id": conversation_id, }, ) - def _expand_messages_for_classification(self, messages: list[dict[str, Any]]) -> list[dict[str, object]]: + def _expand_messages_for_classification(self, messages: Sequence[Mapping[str, object]]) -> list[dict[str, object]]: """ Convert raw OpenAI-format messages to Lasso API format with content blocks. @@ -701,7 +704,7 @@ class LassoGuardrail(CustomGuardrail): expanded.append({"role": role, "content": content}) if role == "assistant": - for call in msg.get("tool_calls") or []: + for call in iter_items(msg.get("tool_calls") or []): call_id, name, input_data = self._extract_tool_call_fields(call) if not call_id or not name: verbose_proxy_logger.warning( @@ -749,7 +752,7 @@ class LassoGuardrail(CustomGuardrail): def _prepare_payload( self, - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT", diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py index 5e2b85250d2..76c3a86d701 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py @@ -13,6 +13,9 @@ import json from pathlib import Path from typing import Any, Final +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base import ( # noqa: F401 # legacy module exports BaseCompetitorIntentChecker, _compile_marker, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -114,6 +117,7 @@ AIRLINE_EXPLICIT_OTHER_MEANING_MARKER: Final = ( ) _MAJOR_AIRLINES_PATH: Final = Path(__file__).resolve().parent / "major_airlines.json" +_DECODED_JSON: Final = TypeAdapter(object) def _load_competitors_excluding_brand(brand_self: list[str]) -> list[str]: @@ -127,11 +131,11 @@ def _load_competitors_excluding_brand(brand_self: list[str]) -> list[str]: return [] try: with open(_MAJOR_AIRLINES_PATH, encoding="utf-8") as f: - airlines: Final = json.load(f) + airlines: Final = _DECODED_JSON.validate_python(json.load(f)) except (json.JSONDecodeError, OSError): return [] result: Final[list[str]] = [] - for entry in airlines: + for entry in iter_items(airlines): if not isinstance(entry, dict): continue match_str = entry.get("match") or "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 40e44151e65..f58c8705815 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -78,6 +78,7 @@ from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1 _DECODED_ARGUMENTS: Final = TypeAdapter(object) +_JSON_ENTRIES: Final = TypeAdapter(list[object]) GAP_WORD_TOKENIZER: Final = re.compile(r"\b\w+\b") SENTENCE_TERMINATORS: Final = re.compile(r"[.!?]+") @@ -699,9 +700,8 @@ class ContentFilterGuardrail(CustomGuardrail): Severity mapping: 4,3 -> high; 2 -> medium; 1 -> low. """ with open(file_path, "r") as f: - entries = json.load(f) - if not isinstance(entries, list): - entries = [entries] + decoded: Final = _DECODED_ARGUMENTS.validate_python(json.load(f)) + entries: Final = _JSON_ENTRIES.validate_python(decoded) if isinstance(decoded, list) else [decoded] # Derive category name from filename (e.g. harm_toxic_abuse.json -> harm_toxic_abuse) category_name: Final = os.path.splitext(os.path.basename(file_path))[0] severity_map: Final = {4: "high", 3: "high", 2: "medium", 1: "low"} @@ -806,7 +806,7 @@ class ContentFilterGuardrail(CustomGuardrail): """ try: with open(file_path, "r") as f: - data: Final = yaml.safe_load(f) + data: Final = _DECODED_ARGUMENTS.validate_python(yaml.safe_load(f)) if not isinstance(data, dict) or "blocked_words" not in data: raise ValueError("Invalid format: file must contain 'blocked_words' key with list of words") diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index 06ea3ce4b20..d2dafb04fe0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -5,7 +5,7 @@ import uuid from collections.abc import AsyncGenerator, AsyncIterable, Iterator, Mapping, Sequence from typing import ( TYPE_CHECKING, - Any, # noqa: TID251 # the streaming hook and the flush copy pass provider chunks through untyped + Any, # noqa: TID251 # the streaming hook passes provider chunks through untyped ClassVar, Final, Literal, @@ -23,7 +23,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) -from litellm.litellm_core_utils.core_helpers import attribute_items +from litellm.litellm_core_utils.core_helpers import attribute_items, attribute_or_none, iter_items, require_attributes from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -99,13 +99,8 @@ class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): pass -class _ContentDelta(Protocol): - content: object - - -class _FinishingDelta(Protocol): - content: object - tool_calls: object +class _DeepCopyable(Protocol): + def model_copy(self, *, deep: bool) -> object: ... class LLMShieldProxyGuardrail(CustomGuardrail): @@ -494,7 +489,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): sse: Final = AnthropicSSERestorer(step) events: Final = ResponsesStreamRestorer(step, rehydrate) carries: Final[CarryWindows] = {} - last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush. + last_chunk = None # rebind-ok: tracks the most recent chunk with choices for the final flush. async for chunk in response: if isinstance(chunk, (bytes, str)): @@ -506,7 +501,8 @@ class LLMShieldProxyGuardrail(CustomGuardrail): yield event continue restored_chunk = detached(chunk) - last_chunk = restored_chunk + if getattr(restored_chunk, "choices", None): + last_chunk = restored_chunk for choice in attribute_items(restored_chunk, "choices"): await self._restore_choice(choice, carries, session_id) yield restored_chunk @@ -527,7 +523,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): one tool call's accumulating `arguments`. A shared window would splice the text held back for one stream onto another. """ - delta: Final = getattr(choice, "delta", None) + delta: Final = attribute_or_none(choice, "delta") index: Final = choice_index(choice) is_final: Final = bool(getattr(choice, "finish_reason", None)) if isinstance(choice, TextChoices): @@ -565,7 +561,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): async def _restore_content_window( self, - delta: _ContentDelta, + delta: object, key: CarryKey, carries: CarryWindows, session_id: str, @@ -573,19 +569,19 @@ class LLMShieldProxyGuardrail(CustomGuardrail): ) -> None: """Restores one delta's content through its own window.""" carry: Final = carries.get(key, "") - text: Final = getattr(delta, "content", None) + text: Final = attribute_or_none(delta, "content") if not isinstance(text, str) or not text: if is_final and carry: flushed, remaining = await self._stream_step("", carry, True, session_id) carries[key] = remaining # rebind-ok: this stream's window advances. if flushed: - delta.content = flushed + write_field(delta, "content", flushed) return emitted, remaining = await self._stream_step(text, carry, is_final, session_id) carries[key] = remaining # rebind-ok: this stream's window advances. - delta.content = emitted + write_field(delta, "content", emitted) async def _restore_tool_call_window( self, @@ -617,7 +613,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): async def _flush_finished_choice( self, - delta: _FinishingDelta, + delta: object, choice_index: int, carries: CarryWindows, session_id: str, @@ -642,12 +638,12 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if not text: continue if tool_index is None: - delta.content = text + write_field(delta, "content", text) else: continuations.extend(continuation_delta(tool_index, text)) if continuations: - existing: Final = tuple(getattr(delta, "tool_calls", None) or ()) - delta.tool_calls = [*existing, *continuations] + existing: Final = tuple(attribute_items(delta, "tool_calls")) + write_field(delta, "tool_calls", [*existing, *continuations]) async def _flush_trailing( self, last_chunk: object, carries: CarryWindows, session_id: str @@ -675,7 +671,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): chunk = self._chunk_for_choice(last_chunk, choice_index) if chunk is None: continue - choice: object = chunk.choices[0] + choice = read_list(chunk, "choices")[0] delta = read_field(choice, "delta") if isinstance(choice, TextChoices): choice.text = text @@ -687,26 +683,26 @@ class LLMShieldProxyGuardrail(CustomGuardrail): yield chunk @staticmethod - def _chunk_for_choice(last_chunk: Any, index: int) -> Any: + def _chunk_for_choice(last_chunk: object, index: int) -> object | None: """A single-choice copy of the last chunk, carrying only `index`. Emitting one choice per chunk keeps a flush from reading as content on a choice it does not belong to. """ - chunk: Final = last_chunk.model_copy(deep=True) - raw_choices: Final = getattr(chunk, "choices", None) + chunk: Final = require_attributes(last_chunk, _DeepCopyable, "model_copy").model_copy(deep=True) + raw_choices: Final = attribute_or_none(chunk, "choices") if not raw_choices: return None - choices: Final[tuple[object, ...]] = tuple(raw_choices) + choices: Final = tuple(iter_items(raw_choices)) position: Final = next((at for at, choice in enumerate(choices) if choice_index(choice) == index), 0) - kept: Final = raw_choices[position] - if getattr(kept, "delta", None) is None and not isinstance(kept, TextChoices): + kept: Final = choices[position] + if attribute_or_none(kept, "delta") is None and not isinstance(kept, TextChoices): return None - kept.index = index - kept.finish_reason = None - chunk.choices = [kept] + write_field(kept, "index", index) + write_field(kept, "finish_reason", None) + write_field(chunk, "choices", [kept]) if hasattr(chunk, "usage"): - del chunk.usage + delattr(chunk, "usage") return chunk async def _stream_step(self, text: str, carry: str, final: bool, session_id: str) -> tuple[str, str]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py index 4ef8bbdbb3e..ef007155c46 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py @@ -7,9 +7,12 @@ from enum import Enum from types import MappingProxyType from typing import ( Final, + Protocol, TypeAlias, ) +from pydantic import TypeAdapter + from .payload import ( JsonBody, MutableRequest, @@ -24,6 +27,8 @@ from .payload import ( write_field, ) +_OPAQUE_VALUE: Final = TypeAdapter(object) + ANTHROPIC_DELTA_FIELDS: Final = MappingProxyType({"text_delta": "text", "input_json_delta": "partial_json"}) SSE_EVENT_BOUNDARY: Final = re.compile(rb"(\r?\n\r?\n)") @@ -80,6 +85,15 @@ def opens_like_sse(head: bytes) -> bool | None: return False +class _ValuedEnum(Protocol): + @property + def value(self) -> object: ... + + +def _enum_value(member: _ValuedEnum) -> object: + return member.value + + def responses_event_type(chunk: object) -> str | None: """The event type of a Responses API stream event, or None for any other chunk. @@ -90,7 +104,7 @@ def responses_event_type(chunk: object) -> str | None: if isinstance(chunk, (bytes, str)): return None kind: Final = read_field(chunk, "type") - value: Final = kind.value if isinstance(kind, Enum) else kind + value: Final = _enum_value(kind) if isinstance(kind, Enum) else kind return value if isinstance(value, str) and value.startswith("response.") else None @@ -173,7 +187,7 @@ class AnthropicSSERestorer: return block line: Final = lines[data_lines[0]] try: - parsed: Final[object] = json.loads(line[len("data:") :]) + parsed: Final = _OPAQUE_VALUE.validate_python(json.loads(line[len("data:") :])) except ValueError: return block event: Final = as_object(parsed) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py index 95db4bd4f77..fa64d306f22 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py @@ -1,4 +1,6 @@ -from typing import TYPE_CHECKING, Any, Final, cast +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol, cast from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -8,13 +10,22 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams +class _ParamsLookup(Protocol): + def get(self, key: str, /) -> object: ... + + +_NO_PARAMS: Final = MappingProxyType[str, object]({}) + + def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): import litellm # Default to always-on. Only disable if the user explicitly sets default_on: false. # We check the raw guardrail dict because LitellmParams normalizes None → False, # making it impossible to distinguish "not set" from "explicitly false" via litellm_params. - _raw_default_on: Final = cast(dict[str, Any], guardrail).get("litellm_params", {}).get("default_on") + _raw_default_on: Final = ( + cast(Mapping[str, _ParamsLookup], guardrail).get("litellm_params", _NO_PARAMS).get("default_on") + ) _default_on: Final = _raw_default_on is not False _callback: Final = MCPEndUserPermissionGuardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 221c4b3752b..2e3f7120726 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -74,14 +74,14 @@ import os import re import time from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional import httpx import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -90,7 +90,6 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral @@ -98,6 +97,10 @@ if TYPE_CHECKING: from jwt.types import Options +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + pass + + class _OIDCDiscoveryDocument(TypedDict, total=False): jwks_uri: str @@ -279,11 +282,13 @@ class MCPJWTSigner(CustomGuardrail): debug_headers: bool = False, # FR-10: Configurable scopes allowed_scopes: list[str] | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: - kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs - super().__init__(**base_kwargs) + options: Final[_CustomGuardrailOptions] = { + "supported_event_hooks": list(self.get_supported_event_hooks()), + **kwargs, + } + super().__init__(**options) # --- Signing key setup --- key_material: Final = os.environ.get(self.SIGNING_KEY_ENV) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 8482567b6da..3cb20a00e2f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -6,7 +6,7 @@ on the LiteLLM gateway. Blocks or alerts when unregistered servers are found. """ from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException @@ -21,6 +21,9 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( ) from litellm.types.guardrails import GuardrailEventHooks +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + class MCPSecurityGuardrail(CustomGuardrail): def __init__( @@ -40,7 +43,7 @@ class MCPSecurityGuardrail(CustomGuardrail): async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, - cache: Any, + cache: "DualCache", data: dict, call_type: str, ) -> Exception | str | dict | None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py index ed955ac829d..06de4ef0e93 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py @@ -1,5 +1,6 @@ from typing import TYPE_CHECKING, Final +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.types.guardrails import SupportedGuardrailIntegrations from .purview_dlp import MicrosoftPurviewDLPGuardrail @@ -11,8 +12,8 @@ if TYPE_CHECKING: def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): import litellm - tenant_id: Final = getattr(litellm_params, "tenant_id", None) - client_id: Final = getattr(litellm_params, "client_id", None) + tenant_id: Final = attribute_or_none(litellm_params, "tenant_id") + client_id: Final = attribute_or_none(litellm_params, "client_id") # client_secret can be passed via the standard api_key field or as # a dedicated client_secret parameter. diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 718401b6e68..43eb54d1464 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -10,9 +10,12 @@ import json import os from collections.abc import Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, TypeAlias, cast +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional, TypeAlias, cast from urllib.parse import urlparse +from pydantic import TypeAdapter +from typing_extensions import TypedDict, Unpack + from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, @@ -25,7 +28,6 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage -from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.noma import NomaV2StreamingParams from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus @@ -39,6 +41,7 @@ _DEFAULT_API_BASE: Final = "https://api.noma.security/" _AIDR_SCAN_ENDPOINT: Final = "/litellm/guardrail" _INTERVENED_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls") _DEFAULT_API_BASE_HOSTNAME: Final = urlparse(_DEFAULT_API_BASE).hostname +_DECODED_JSON: Final = TypeAdapter(object) _GuardrailJsonResponse: TypeAlias = Exception | str | dict[str, object] @@ -50,6 +53,10 @@ _LOGGING_KEYS_DUPLICATING_SCAN_INPUTS: Final = _KEYS_DUPLICATING_SCAN_INPUTS + ( ) +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + pass + + class _Action(str, enum.Enum): BLOCKED = "BLOCKED" NONE = "NONE" @@ -74,7 +81,7 @@ class NomaV2Guardrail(CustomGuardrail): streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, streaming_buffer_release_on_scan: bool | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -115,10 +122,11 @@ class NomaV2Guardrail(CustomGuardrail): if streaming.streaming_buffer_release_on_scan is not None: self.streaming_buffer_release_on_scan: bool = streaming.streaming_buffer_release_on_scan - kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - - base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs - super().__init__(**base_kwargs) + options: Final[_CustomGuardrailOptions] = { + "supported_event_hooks": list(self.get_supported_event_hooks()), + **kwargs, + } + super().__init__(**options) @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: @@ -219,7 +227,7 @@ class NomaV2Guardrail(CustomGuardrail): except (ValueError, TypeError): json_str = safe_dumps(payload) - safe_payload: Final[object] = safe_json_loads(json_str, default={}) + safe_payload: Final = _DECODED_JSON.validate_python(safe_json_loads(json_str, default={})) if safe_payload == {} and payload: verbose_proxy_logger.warning( "Noma v2 guardrail: payload serialization failed, falling back to empty payload" diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 1a6feb47215..7c26c4c4c56 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Final, Literal, Optional import httpx from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -27,6 +28,8 @@ from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + class OnyxGuardrail(CustomGuardrail): @classmethod @@ -115,7 +118,7 @@ class OnyxGuardrail(CustomGuardrail): try: response: Final = ModelResponse(**request_data) parsed: Final = response.json() - payload = parsed.get("response", {}) + payload = _DECODED_JSON.validate_python(parsed.get("response", {})) except Exception as e: verbose_proxy_logger.error( "Error in converting request_data to ModelResponse: %s", diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py index fe526b72ac3..4d120a2dd26 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING, Final import litellm +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( OpenAIModerationGuardrail, ) @@ -40,12 +41,12 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" return openai_moderation_guardrail -def _get_config_value(litellm_params, optional_params, attribute_name): +def _get_config_value(litellm_params: object, optional_params: object, attribute_name: str) -> object: if optional_params is not None: - value: Final = getattr(optional_params, attribute_name, None) + value: Final = attribute_or_none(optional_params, attribute_name) if value is not None: return value - return getattr(litellm_params, attribute_name, None) + return attribute_or_none(litellm_params, attribute_name) guardrail_initializer_registry: Final = { diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 7af5eaf1da9..98be35079c3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -8,6 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -40,6 +41,8 @@ if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIModerationResponse from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + def _config_secret(value: object) -> object: return get_secret_str(value) if isinstance(value, str) and value.startswith("os.environ/") else value @@ -151,7 +154,9 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): timeout=self.timeout, ) - verbose_proxy_logger.debug("OpenAI Moderation guard response: %s", response.json()) + verbose_proxy_logger.debug( + "OpenAI Moderation guard response: %s", _DECODED_JSON.validate_python(response.json()) + ) if response.status_code != 200: raise HTTPException( diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 8e03f7db4f5..6426bfb8788 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -18,6 +18,7 @@ from urllib.parse import urlparse import httpx from fastapi import HTTPException from pydantic import ConfigDict, TypeAdapter, ValidationError, field_validator +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -129,6 +130,21 @@ _RESPONSES_INPUT: Final[TypeAdapter[_ResponsesInput]] = TypeAdapter(_ResponsesIn _DECODED_ARGUMENTS: Final = TypeAdapter(object) +class _AirsToolEvent(TypedDict): + metadata: ReadOnly[Mapping[str, object]] + input: NotRequired[ReadOnly[str]] + + +def _mcp_tool_input(mcp_arguments: object) -> str | None: + if mcp_arguments is not None and mcp_arguments != "": + serialized_args: Final = ( + json.dumps(mcp_arguments) if isinstance(mcp_arguments, (dict, list)) else str(mcp_arguments) + ) + if serialized_args.strip(): + return serialized_args + return None + + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -339,7 +355,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): is_response: bool = False, metadata: Mapping[str, object] | None = None, call_id: object = None, - tool_event: dict[str, Any] | None = None, + tool_event: _AirsToolEvent | None = None, ) -> dict[str, object]: """Call PANW Prisma AIRS API to scan content or a tool_event.""" @@ -894,7 +910,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): if not line.startswith("data: "): continue try: - data = json.loads(line[6:]) + data = _DECODED_ARGUMENTS.validate_python(json.loads(line[6:])) except (json.JSONDecodeError, ValueError): continue if not isinstance(data, dict): @@ -1954,24 +1970,21 @@ class PanwPrismaAirsHandler(CustomGuardrail): # "mcp_tool_name"/"mcp_arguments". Check canonical first, then fallback. mcp_tool_name: Final = request_data.get("mcp_tool_name") or self._mcp_name_fallback(request_data) if mcp_tool_name and input_type == "request": - mcp_tool_event: Final[dict[str, object]] = { - "metadata": { - "ecosystem": "mcp", - "method": "tools/call", - "server_name": self._get_mcp_server_name(request_data, mcp_tool_name), - "tool_invoked": mcp_tool_name, - }, + mcp_tool_metadata: Final = { + "ecosystem": "mcp", + "method": "tools/call", + "server_name": self._get_mcp_server_name(request_data, mcp_tool_name), + "tool_invoked": mcp_tool_name, } mcp_arguments = request_data.get("mcp_arguments") if mcp_arguments is None: mcp_arguments = request_data.get("arguments") - if mcp_arguments is not None and mcp_arguments != "": - if isinstance(mcp_arguments, (dict, list)): - serialized_args = json.dumps(mcp_arguments) - else: - serialized_args = str(mcp_arguments) - if serialized_args.strip(): - mcp_tool_event["input"] = serialized_args + mcp_tool_input: Final = _mcp_tool_input(mcp_arguments) + mcp_tool_event: Final[_AirsToolEvent] = ( + {"metadata": mcp_tool_metadata} + if mcp_tool_input is None + else {"metadata": mcp_tool_metadata, "input": mcp_tool_input} + ) mcp_scan_result: Final = await self._call_panw_api( tool_event=mcp_tool_event, diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index cadc9ca689b..7e2f6bfec33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -85,6 +85,7 @@ class _PresidioAnonymizeResponse(TypedDict): _PRESIDIO_ANONYMIZE_ADAPTER: Final[TypeAdapter[_PresidioAnonymizeResponse | None]] = TypeAdapter( _PresidioAnonymizeResponse | None ) +_DECODED_JSON: Final = TypeAdapter(object) class _JsonResponse(Protocol): @@ -1595,10 +1596,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if line.startswith("data: ") and line != "data: [DONE]": raw_json = line[6:] try: - event = json.loads(raw_json) + event = _DECODED_JSON.validate_python(json.loads(raw_json)) delta = event.get("delta") if isinstance(event, dict) else None if ( - isinstance(delta, dict) + isinstance(event, dict) + and isinstance(delta, dict) and event.get("type") == "content_block_delta" and delta.get("type") == "text_delta" and isinstance(delta.get("text"), str) diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index a92e67046bd..6150fccac0c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -10,6 +10,7 @@ import os from typing import Any, Final, Literal from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -29,6 +30,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs GUARDRAIL_NAME: Final = "qualifire" DEFAULT_QUALIFIRE_API_BASE: Final = "https://proxy.qualifire.ai" +_DECODED_JSON: Final = TypeAdapter(object) class QualifireGuardrail(CustomGuardrail): @@ -188,7 +190,7 @@ class QualifireGuardrail(CustomGuardrail): args = function_info.get("arguments", {}) if isinstance(args, str): try: - args = json.loads(args) + args = _DECODED_JSON.validate_python(json.loads(args)) except json.JSONDecodeError: args = {} api_tool_calls.append( diff --git a/litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py new file mode 100644 index 00000000000..a15a372c462 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py @@ -0,0 +1,72 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, TypeAdapter + +from litellm.secret_managers.main import get_secret_str +from litellm.types.guardrails import SupportedGuardrailIntegrations +from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoOptionalParams + +from .reco import RecoGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + +_CONFIG_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def _optional_params_mapping(optional_params: object) -> Mapping[str, object]: + """optional_params arrives as a pydantic object from YAML, or as a plain dict when reconciled from the DB.""" + if optional_params is None: + return MappingProxyType({}) + raw: Final = ( + optional_params.model_dump(exclude_unset=True) if isinstance(optional_params, BaseModel) else optional_params + ) + return _CONFIG_MAPPING.validate_python(raw) + + +def _first_configured(field_name: str, sources: tuple[Mapping[str, object], ...]) -> object: + value: Final = next((source[field_name] for source in sources if source.get(field_name) is not None), None) + if isinstance(value, str) and value.startswith("os.environ/"): + return get_secret_str(value) + return value + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RecoGuardrail: + import litellm + + # Top-level fields only count when set explicitly: their defaults are fail-closed and would flip Reco's. + explicit_top_level: Final = _CONFIG_MAPPING.validate_python( + litellm_params.model_dump(include=litellm_params.model_fields_set & RecoOptionalParams.model_fields.keys()) + ) + sources: Final = (_optional_params_mapping(litellm_params.optional_params), explicit_top_level) + configured: Final = ( + (field_name, _first_configured(field_name, sources)) for field_name in RecoOptionalParams.model_fields + ) + reco_params: Final = RecoOptionalParams.model_validate( + {field_name: value for field_name, value in configured if value is not None} + ) + + _reco_callback: Final = RecoGuardrail( + guardrail_name=guardrail["guardrail_name"], + reco_tenant_id=reco_params.reco_tenant_id, + api_base=reco_params.api_base, + unreachable_fallback=reco_params.unreachable_fallback, + fail_on_error=reco_params.fail_on_error, + event_hook=litellm_params.mode, + default_on=bool(litellm_params.default_on), + ) + litellm.logging_callback_manager.add_litellm_callback(_reco_callback) + + return _reco_callback + + +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.RECO.value: initialize_guardrail, +} + + +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.RECO.value: RecoGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/reco/reco.py b/litellm/proxy/guardrails/guardrail_hooks/reco/reco.py new file mode 100644 index 00000000000..c8998a20367 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/reco/reco.py @@ -0,0 +1,84 @@ +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final, Literal, Optional + +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import ( + log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped +) +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( + GenericGuardrailAPI, +) +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.proxy.guardrails.guardrail_hooks.reco import validate_reco_tenant_id +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoOptionalParams + + +class RecoGuardrail(GenericGuardrailAPI): + """Reco guardrail integration for LiteLLM, built on the Generic Guardrail API wire contract.""" + + def __init__( + self, + reco_tenant_id: str | None, + api_base: str | None, + headers: Mapping[str, str] | None = None, + unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", + fail_on_error: bool = True, + guardrail_name: str | None = None, + event_hook: str | Sequence[str] | Mode | None = None, + default_on: bool = False, + ) -> None: + if not reco_tenant_id: + raise ValueError("reco_tenant_id is required for the Reco guardrail") + if not api_base: + raise ValueError("api_base is required for the Reco guardrail") + + validated_tenant_id: Final = validate_reco_tenant_id(reco_tenant_id) + + super().__init__( # pyright: ignore[reportUnknownMemberType] # parent takes untyped extra_headers and **kwargs + headers={**(headers or {}), "X-Reco-Tenant-Id": validated_tenant_id}, + api_base=api_base, + unreachable_fallback=unreachable_fallback, + fail_on_error=fail_on_error, + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + ) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], # mutable-ok: must match GenericGuardrailAPI.apply_guardrail's own param type + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + # GenericGuardrailAPI raises with a hardcoded module-level guardrail_name on BLOCKED, + # not the guardrail's own configured name. Re-raise with the right one instead of + # patching the shared base class. + try: + return await super().apply_guardrail( # pyright: ignore[reportUnknownMemberType] # parent annotates request_data as a bare dict + inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj + ) + except GuardrailRaisedException as e: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=e.message, + should_wrap_with_default_message=False, + status_code=e.status_code, + blocked_content=e.blocked_content, + ) from e + + @staticmethod + def get_config_model() -> type["GuardrailConfigModel[RecoOptionalParams]"] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoConfigModel + + return RecoConfigModel + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: fixed by base class + return [GuardrailEventHooks.pre_call] diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py index 1a571047977..342e5b64714 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py @@ -10,6 +10,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final import yaml +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD @@ -22,6 +23,7 @@ if TYPE_CHECKING: ROUTE_TEMPLATES_DIR: Final = os.path.join(os.path.dirname(__file__), "route_templates") +_DECODED_YAML: Final = TypeAdapter(object) class SemanticGuardRouteLoader: @@ -55,7 +57,7 @@ class SemanticGuardRouteLoader: if not os.path.exists(file_path): raise ValueError(f"SemanticGuard: custom routes file not found: {file_path}") with open(file_path, "r") as f: - data: Final = yaml.safe_load(f) + data: Final = _DECODED_YAML.validate_python(yaml.safe_load(f)) if isinstance(data, list): return data if isinstance(data, dict): diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py index 2fbd50b5863..37acee28b53 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py @@ -6,13 +6,14 @@ via embedding similarity. Smarter than regex (understands intent), lighter than an LLM call (~20-50ms per request for embedding). """ -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from litellm._logging import verbose_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import attribute_items, attribute_or_none, require_attributes from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.route_loader import ( SemanticGuardRouteLoader, ) @@ -188,17 +189,21 @@ def _extract_user_text(messages: list) -> str: return "" -def _extract_response_text(response: Any) -> str: +class _MessageContent(Protocol): + @property + def content(self) -> object: ... + + +def _extract_response_text(response: object) -> str: """Extract text from every LLM response choice.""" - if hasattr(response, "choices") and response.choices: - text_parts: Final[list[str]] = [] - for choice in response.choices: - if hasattr(choice, "message") and choice.message: - text = _content_to_text(choice.message.content) - if text: - text_parts.append(text) - return "\n".join(text_parts) - return "" + text_parts: Final[list[str]] = [] + for choice in attribute_items(response, "choices"): + message = attribute_or_none(choice, "message") + if message: + text = _content_to_text(require_attributes(message, _MessageContent, "content").content) + if text: + text_parts.append(text) + return "\n".join(text_parts) def _content_to_text(content: object) -> str: diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index de0f8856522..fbf1f2576ff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -105,7 +105,7 @@ def _resolve_tool_names_from_messages(messages: list[dict]) -> dict[str, str]: fn = (tc.get("function") or {}).get("name") else: tc_id = getattr(tc, "id", None) - fn_obj = getattr(tc, "function", None) + fn_obj: object = getattr(tc, "function", None) fn = getattr(fn_obj, "name", None) if fn_obj else None if tc_id and fn: mapping[tc_id] = fn 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 3ad9d1a2c42..d713f1078a2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime import anyio from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -50,6 +51,8 @@ A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME: Final = "unified_llm_guardrails" +_DECODED_JSON: Final = TypeAdapter(object) + _RequestData: TypeAlias = dict[str, object] @@ -183,7 +186,7 @@ def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) return item.get("id") if isinstance(item, str): try: - obj = json.loads(item.strip()) + obj = _DECODED_JSON.validate_python(json.loads(item.strip())) if isinstance(obj, dict) and "id" in obj: return obj.get("id") except (json.JSONDecodeError, TypeError): @@ -729,7 +732,7 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict: UserAPIKeyAuth, call_type: str, sampling_rate: int, - end_of_stream_only: bool, + end_of_stream_only: object, mappings: Mapping[CallTypes, type["BaseTranslation"]], ) -> AsyncGenerator[object, None]: """Emit guardrail text transformations as new deltas on the stream. @@ -1135,16 +1138,19 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply = request_data.pop("guardrail_to_apply", None) typed_request_data: Final[_RequestData] = request_data - def _streaming_flag(name: str, default: object) -> Any: + def _streaming_setting(name: str, default: object) -> object: return self.resolve_streaming_flag(guardrail_to_apply, name, default) + def _streaming_flag(name: str, default: object) -> Any: + return _streaming_setting(name, default) + sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). - end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False) + end_of_stream_only: object = _streaming_setting("streaming_end_of_stream_only", False) # "block_only" (default) drops guardrail text rewrites on the streaming # path; "incremental_diff" emits them as synthetic deltas (see # _run_incremental_transform_stream). - streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only") + streaming_transform_mode: Final = _streaming_setting("streaming_transform_mode", "block_only") # Withhold every chunk until end-of-stream moderation passes, then # release the original chunks (clean) or only the block message # (blocked) -- moderating the whole response *before* any content @@ -1152,10 +1158,10 @@ class UnifiedLLMGuardrails(CustomLogger): # release the original chunks are replayed as-is, so a # content-rewriting guardrail (e.g. PII masking) would leak # unredacted content. Guarded below via mask_response_content. - buffer_until_moderated: bool = _streaming_flag( + buffer_until_moderated: object = _streaming_setting( "streaming_buffer_until_moderated", buffer_until_moderated_default ) - release_on_scan: Final[bool] = _streaming_flag("streaming_buffer_release_on_scan", False) + release_on_scan: Final = _streaming_setting("streaming_buffer_release_on_scan", False) if ( buffer_until_moderated diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 7521c1e9f45..24b30e65e2b 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -102,6 +102,16 @@ def _as_json_array(value: object) -> list[object] | None: return _JSON_ARRAY.validate_python(value) if isinstance(value, list) else None +def _extra_param_items(litellm_params: LitellmParams) -> Iterable[tuple[str, object]]: + return ( + litellm_params.model_dump(exclude_none=True).items() + if hasattr(litellm_params, "model_dump") + else iter(litellm_params) + if litellm_params + else () + ) + + def contains_encrypted_marker(value: object, depth: int = 0) -> bool: """True if any string in value, at any JSON depth, starts with the encrypted-value prefix.""" if depth > DEFAULT_MAX_RECURSE_DEPTH: @@ -904,13 +914,7 @@ class InMemoryGuardrailHandler: # This matches the behavior of other guardrail initializers (e.g., initialize_lakera) # and aligns with the documented behavior for custom guardrails excluded_extra_param_keys: Final = frozenset(("guardrail", "mode", "default_on", "stream_scope")) - extra_params_items: Final = ( - litellm_params.model_dump(exclude_none=True).items() - if hasattr(litellm_params, "model_dump") - else iter(litellm_params) - if litellm_params - else () - ) + extra_params_items: Final = _extra_param_items(litellm_params) extra_params: Final = MappingProxyType( {key: value for key, value in extra_params_items if key not in excluded_extra_param_keys} ) diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py index 526f58dcd36..c49835905fb 100644 --- a/litellm/proxy/guardrails/tool_name_extraction.py +++ b/litellm/proxy/guardrails/tool_name_extraction.py @@ -6,8 +6,8 @@ knowledge lives in one place. Uses guardrail translation handlers where availabl with standalone extractors for generate_content and MCP. """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Callable, Mapping +from typing import Final from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings @@ -15,7 +15,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import list_field from litellm.types.utils import CallTypes # Call types that have no guardrail translation handler; we use standalone extractors -STANDALONE_EXTRACTORS: Final[dict[str, Any]] = {} +STANDALONE_EXTRACTORS: Final[dict[str, Callable[[Mapping[str, object]], list[str]]]] = {} def _extract_generate_content_tool_names(data: Mapping[str, object]) -> list[str]: @@ -28,7 +28,7 @@ def _extract_generate_content_tool_names(data: Mapping[str, object]) -> list[str return names -def _extract_mcp_tool_names(data: dict) -> list[str]: +def _extract_mcp_tool_names(data: Mapping[str, object]) -> list[str]: """MCP call_tool: name or mcp_tool_name in body""" names: Final[list[str]] = [] name: Final = data.get("name") or data.get("mcp_tool_name") diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 2f466f22ef8..caa66bddf39 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -11,10 +11,11 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload from fastapi import APIRouter, Depends, Query -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.usage_tracking import guardrail_status_to_action @@ -838,28 +839,41 @@ def _snippet(text: object, max_len: int = 200) -> str | None: return result -def _input_snippet_for_log(sl: "prisma_models.LiteLLM_SpendLogs") -> str | None: - """Snippet for request input: prefer messages, fall back to proxy_server_request (same as drawer).""" - out = _snippet(sl.messages) +_DECODED_JSON: Final = TypeAdapter(object) + + +def _request_dict(value: object) -> Mapping[str, object] | None: + mapping: Final = as_str_mapping(value) + return mapping if isinstance(mapping, dict) else None + + +def _proxy_request_snippet(psr: object) -> str | None: + request: Final = _request_dict(psr) + if request is None: + return _snippet(psr) + msgs: Final = request.get("messages") + body: Final = _request_dict(request.get("body")) if msgs is None else None + out: Final = _snippet(body.get("messages") if body is not None else msgs) if out: return out - psr = getattr(sl, "proxy_server_request", None) + return _snippet(psr) + + +def _input_snippet_for_log(sl: "prisma_models.LiteLLM_SpendLogs") -> str | None: + """Snippet for request input: prefer messages, fall back to proxy_server_request (same as drawer).""" + out: Final = _snippet(sl.messages) + if out: + return out + psr: Final[object] = getattr(sl, "proxy_server_request", None) if not psr: return None - if isinstance(psr, str): - try: - psr = json.loads(psr) - except Exception: - return _snippet(psr) - if isinstance(psr, dict): - msgs = psr.get("messages") - if msgs is None and isinstance(psr.get("body"), dict): - msgs = psr["body"].get("messages") - out = _snippet(msgs) - if out: - return out + if not isinstance(psr, str): + return _proxy_request_snippet(psr) + try: + decoded: Final = _DECODED_JSON.validate_python(json.loads(psr)) + except Exception: return _snippet(psr) - return _snippet(psr) + return _proxy_request_snippet(decoded) @router.get( diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 684f1619512..53c54fe15bf 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -2304,7 +2304,15 @@ async def test_model_connection( ) selected_by_id: Final = deployment_by_id is not None - reject_server_owned_wif_fields(request_litellm_params) + reject_server_owned_wif_fields( + MappingProxyType( + { + key: value + for key, value in _OBJECT_MAPPING.validate_python(request_litellm_params).items() + if key != "chatgpt_auth_file" + } + ) + ) # Merge: config params (from proxy config) as base, request params override litellm_params = { **_config_base_for_health_check( diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 47c3869fb64..8f9aebecf30 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -455,10 +455,10 @@ class KeyManagementEventHooks: try: from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.proxy_server import prisma_client except ImportError: return None - prisma_client: Final = getattr(proxy_server_module, "prisma_client", None) user_api_key_cache: Final = getattr(proxy_server_module, "user_api_key_cache", None) if prisma_client is None or user_api_key_cache is None: diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 37cef41034b..bc9cee1cb1d 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -374,6 +374,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations = None + self._retired = False + self._flush_after_retire = None self.deployment_budget_config = None @with_service_target("model_budgets") diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 47ea50debc9..05c64317a44 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -892,7 +892,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): {**_hidden_params, "additional_headers": _additional_headers}, ) - return await super().async_post_call_success_hook(data, user_api_key_dict, response) + await super().async_post_call_success_hook(data, user_api_key_dict, response) PROXY_MaxParallelRequestsHandler: Final = _PROXY_MaxParallelRequestsHandler diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 4f2f7f9c256..567b8452ae2 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -73,6 +73,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction from litellm.router_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, @@ -100,7 +101,7 @@ if TYPE_CHECKING: from litellm.types.agents import AgentResponse from litellm.types.caching import RedisPipelineIncrementOperation - Span = _Span | Any + Span = _Span InternalUsageCache = _InternalUsageCache else: Span = Any @@ -695,6 +696,7 @@ class RequestRateLimiterStash: owner_litellm_call_id: str | None = None rate_limit_response: RateLimitResponse | None = None parallel_slot: ParallelSlotAcquisition | None = None + parallel_slot_reused_at: datetime | None = None parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) reserved_tokens: int = 0 reserved_model: RateLimitedModel | None = None @@ -791,6 +793,19 @@ def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: return stash +def _ended_before_slot_reuse(ended_at: datetime | None, reused_at: datetime | None) -> bool: + if ended_at is None or reused_at is None or (ended_at.tzinfo is None) != (reused_at.tzinfo is None): + return False + return ended_at <= reused_at + + +def _callback_kwargs_mapping(kwargs: object) -> Mapping[str, object]: + try: + return _REQUEST_RATE_LIMIT_DATA.validate_python(kwargs) + except ValidationError: + return MappingProxyType({}) + + def get_request_stash_for_call(litellm_call_id: str | None) -> RequestRateLimiterStash | None: stash: Final = _request_stash.get() if stash is None: @@ -1233,7 +1248,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): estimable_data.get("input"), ) ) - messages, prompt, input_text = selected_fields + messages, prompt, raw_input = selected_fields + input_text: Final = ( + billable_input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input + ) total_chars: Final = ( len(get_str_from_messages(messages)) @@ -1652,6 +1670,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): read_only: bool = False, skip_tpm_check: bool = False, parallel_slot_id: str | None = None, + skip_parallel_check: bool = False, ) -> RateLimitResponse: """ Check if any of the rate limit descriptors should be rate limited. @@ -1668,6 +1687,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ``reserve_tpm_tokens`` reservation path should set this to avoid the +1-per-key Lua / in-memory increment double-charging the tokens counter. + skip_parallel_check: If True, ``max_parallel_requests`` gauges are + neither checked nor acquired, for a call that runs inside a + request already holding its slot. ``max_parallel_requests`` descriptors are enforced by the dedicated concurrency-gauge path (``_check_parallel_request_gauges``), never by @@ -1760,7 +1782,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if windowed_response["overall_code"] == "OVER_LIMIT": return windowed_response - if not gauges: + if not gauges or skip_parallel_check: return windowed_response gauge_response: Final = await self._check_parallel_request_gauges( @@ -2023,12 +2045,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span: Span | None, *, in_logging_callback: bool = False, + round_ended_at: datetime | None = None, ) -> None: if stash is None: return async with stash.parallel_slot_release_lock: acquisition: Final = stash.parallel_slot - if acquisition is None: + if acquisition is None or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at): return deferred: Final = in_logging_callback and await self._defer_parallel_slot_release( acquisition, parent_otel_span @@ -3816,6 +3839,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=data.get("input") or "", responses_api_request=data, + only_since_compaction=True, ) @staticmethod @@ -3947,6 +3971,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_reservation_amount: int, call_type: str | None = None, endpoint_type: EndpointType = EndpointType.GENERIC, + reuses_request_slot: bool = False, ) -> None: """ Reserve project-scoped ITPM/OTPM tokens (Bedrock Mantle-style @@ -4026,7 +4051,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.reservation_released = True - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=io_response, descriptors=descriptors, @@ -4187,6 +4213,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") stash: Final = claim_request_stash_for_data(data) + reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None + if reuses_request_slot: + stash.parallel_slot_reused_at = self._get_current_time() ######################################################### # Check if the call type has a specific rate limiter @@ -4231,11 +4260,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # then because _reserve_project_io_tokens_or_raise below charges # them unconditionally and counting them here too would # double-charge every request. - parallel_counter_keys: Final = [ - self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") - for d in descriptors - if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None - ] + parallel_counter_keys: Final = ( + [] + if reuses_request_slot + else [ + self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") + for d in descriptors + if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None + ] + ) parallel_slot_id: Final = uuid.uuid4().hex if parallel_counter_keys else None first_pass_descriptors: Final = ( @@ -4250,6 +4283,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, skip_tpm_check=self.tpm_reservation_enabled, parallel_slot_id=parallel_slot_id, + skip_parallel_check=reuses_request_slot, ) if response["overall_code"] == "OVER_LIMIT": @@ -4259,7 +4293,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, ) else: - stash.rate_limit_response = response + if not reuses_request_slot: + stash.rate_limit_response = response if parallel_slot_id is not None: stash.parallel_slot = ParallelSlotAcquisition( slot_id=parallel_slot_id, @@ -4357,7 +4392,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): raw_estimate=estimated_tokens, ) if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit: - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) raise PTUShareExceededError(stash.ptu_ceiling, ptu_estimated_tokens) tpm_response: Final = await self.reserve_tpm_tokens( descriptors=descriptors, @@ -4367,7 +4403,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if tpm_response["overall_code"] == "OVER_LIMIT": - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, @@ -4418,6 +4455,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_reservation_amount=tpm_reservation_amount, call_type=call_type, endpoint_type=endpoint_type, + reuses_request_slot=reuses_request_slot, ) def _create_pipeline_operations( @@ -5332,8 +5370,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): try: verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") - stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True) + callback_kwargs: Final = _callback_kwargs_mapping(kwargs) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(callback_kwargs)) + await self._release_stashed_parallel_slot( + stash, + litellm_parent_otel_span, + in_logging_callback=True, + round_ended_at=end_time + if isinstance(end_time, datetime) and not callback_kwargs.get("stream") + else None, + ) pipeline_operations: Final = self._build_success_event_pipeline_operations( kwargs=kwargs, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 2062d154a75..e784ec90361 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost +from litellm.litellm_core_utils.redact_messages import MESSAGE_REDACTION_ENABLE_HEADERS from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( @@ -38,6 +39,7 @@ from litellm.proxy.db.db_spend_update_writer import ( ) from litellm.proxy.hooks.search_free_tier import apply_search_free_tier from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route +from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope from litellm.proxy.spend_tracking.spend_event import ( ObjectMapping, @@ -60,7 +62,12 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # sanitize_error_information_for_spend_logs, should_store_prompts_and_responses_in_spend_logs, ) -from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend +from litellm.proxy.utils import ( + REJECTED_BEFORE_REQUEST_SETUP, + REJECTED_BEFORE_REQUEST_SETUP_KEY, + BilledFailedAttempt, + ProxyUpdateSpend, +) from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, @@ -130,6 +137,9 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( FAILURE_SPEND_COUNTER_WAIT_SECONDS: Final = 1.0 +_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) + + def _proxy_stamped_used_client_oauth_token( request_data: Mapping[str, object], request_route: str | None ) -> bool | None: @@ -186,6 +196,26 @@ def _failure_snapshot_with_guardrails_recorded_after_it( return {**failure_snapshot, "guardrail_information": recorded_guardrail_information} +def _guardrail_info_lifted_from_litellm_metadata( + litellm_metadata: ObjectMapping, *, guardrail_info_already_set: bool +) -> ObjectMapping: + guardrail_info: Final = litellm_metadata.get("standard_logging_guardrail_information") + if guardrail_info is None or guardrail_info_already_set: + return {} + return {"standard_logging_guardrail_information": guardrail_info} + + +def _redaction_opt_in_lifted_from_litellm_metadata( + litellm_metadata: ObjectMapping, *, caller_set_headers: bool +) -> ObjectMapping: + request_headers: Final = litellm_metadata.get("headers") + if caller_set_headers or not isinstance(request_headers, dict): + return {} + headers: Final = _OBJECT_MAPPING.validate_python(request_headers) + opt_in_headers: Final = {name: headers[name] for name in MESSAGE_REDACTION_ENABLE_HEADERS if name in headers} + return {"headers": opt_in_headers} if opt_in_headers else {} + + def _proxy_spend_writer() -> DBSpendUpdateWriter: from litellm.proxy.proxy_server import proxy_logging_obj @@ -386,6 +416,8 @@ class _ProxyDBLogger(CustomLogger): request_route: Final = user_api_key_dict.request_route if ( ProxyDBLogger._should_track_errors_in_db() is False + or ProxyDBLogger._skips_pre_routing_rejections() + and is_pre_routing_rejection(original_exception, _REQUEST_DATA.validate_python(request_data)) or request_route is not None and not ( RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route) @@ -431,13 +463,17 @@ class _ProxyDBLogger(CustomLogger): existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata}) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") - if ( - isinstance(litellm_metadata_bucket, dict) - and "standard_logging_guardrail_information" not in existing_metadata - ): - guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information") - if guardrail_info is not None: - existing_metadata["standard_logging_guardrail_information"] = guardrail_info + litellm_metadata: Final[ObjectMapping] = ( + _OBJECT_MAPPING.validate_python(litellm_metadata_bucket) + if isinstance(litellm_metadata_bucket, dict) + else {} + ) + existing_metadata.update( + _guardrail_info_lifted_from_litellm_metadata( + litellm_metadata, + guardrail_info_already_set="standard_logging_guardrail_information" in existing_metadata, + ) + ) if "litellm_params" not in request_data: request_data["litellm_params"] = {} @@ -452,7 +488,12 @@ class _ProxyDBLogger(CustomLogger): request_data["litellm_params"]["proxy_server_request"] = ( request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} ) - request_data["litellm_params"]["metadata"] = existing_metadata + request_data["litellm_params"]["metadata"] = { + **existing_metadata, + **_redaction_opt_in_lifted_from_litellm_metadata( + litellm_metadata, caller_set_headers="headers" in existing_metadata + ), + } # Preserve model name and custom_llm_provider if "model" not in request_data: @@ -931,16 +972,40 @@ class _ProxyDBLogger(CustomLogger): If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings """ - from litellm.proxy.proxy_server import general_settings - - if general_settings.get("disable_error_logs") is True: + if _general_setting_enabled("disable_error_logs"): return False return + @staticmethod + def _skips_pre_routing_rejections() -> bool: + """``general_settings.disable_pre_routing_error_logs`` keeps requests the proxy rejected before + selecting a deployment out of the spend logs, so a client retrying a bad key or an exhausted + budget cannot flood ``LiteLLM_SpendLogs`` with zero-spend rows.""" + return _general_setting_enabled("disable_pre_routing_error_logs") + + +def _general_setting_enabled(name: str) -> bool: + from litellm.proxy.proxy_server import general_settings + + return name in general_settings and general_settings[name] is True + ProxyDBLogger: Final = _ProxyDBLogger +_REQUEST_DATA: Final = TypeAdapter(dict[str, object]) + + +def is_pre_routing_rejection(exception: Exception, request_data: Mapping[str, object]) -> bool: + """The pipeline sets up the logging object only after auth, body parsing and model resolution succeed, and + before guardrails or routing run. ``ProxyLogging._handle_logging_proxy_only_error`` marks a failure that + arrives without one with the ``REJECTED_BEFORE_REQUEST_SETUP`` sentinel, so that request never had a provider + considered. An unknown model is the one later rejection that still never reached a provider.""" + return request_data.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP or isinstance( + exception, ProxyModelNotFoundError + ) + + def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: patch = {k: v for k, v in metadata.items() if (k.startswith("user_api_key") or k == "tags") and v is not None} if not patch: diff --git a/litellm/proxy/hooks/rate_limiter_utils.py b/litellm/proxy/hooks/rate_limiter_utils.py index bbe38068638..c4949834896 100644 --- a/litellm/proxy/hooks/rate_limiter_utils.py +++ b/litellm/proxy/hooks/rate_limiter_utils.py @@ -7,6 +7,7 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.types.router import ModelGroupInfo from litellm.types.utils import PriorityReservationDict @@ -85,10 +86,10 @@ def _resolve_provider_from_router_alias( if llm_router is None: return None try: - model_list: Final = getattr(llm_router, "model_list", None) + model_list: Final[object] = getattr(llm_router, "model_list", None) if not model_list: return None - for deployment in model_list: + for deployment in iter_items(model_list): if not isinstance(deployment, dict): continue if deployment.get("model_name") != model: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 09f53cecb72..b51eea7a876 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -268,7 +268,7 @@ def sanitize_for_log(value: object) -> str: from litellm.router import Router from litellm.secret_managers.main import get_secret_bool -from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES +from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES, STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.router import BaggageMappedTagsStamp, holds_api_key_secret_reference from litellm.types.services import ServiceTypes @@ -334,6 +334,7 @@ LITELLM_METADATA_ROUTES: Final = ( "/v1/messages", "responses", "files", + "interactions", ) LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( @@ -353,6 +354,7 @@ LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "weights", "_router_weights", + "_encrypted_content_affinity_pinned", "fallback_depth", "_target_order", "attempted_targets", @@ -1415,13 +1417,14 @@ def clean_headers( Cleaned headers dict """ from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + from litellm.llms.chatgpt.common_utils import is_chatgpt_oauth_key clean_headers: Final = {} litellm_key_lower: Final = litellm_key_header_name.lower() if litellm_key_header_name is not None else None for header, value in headers.items(): header_lower = header.lower() - if header_lower == "authorization" and is_anthropic_oauth_key(value): + if header_lower == "authorization" and (is_anthropic_oauth_key(value) or is_chatgpt_oauth_key(value)): if authenticated_with_header is None or authenticated_with_header.lower() != "authorization": clean_headers[header] = value continue @@ -1450,7 +1453,7 @@ def clean_headers( def _is_credential_header(header: str) -> bool: """Whether `header` carries a caller credential rather than request context.""" - return header.lower() in _CREDENTIAL_HEADER_NAMES + return header.lower().removeprefix("x-pass-") in _CREDENTIAL_HEADER_NAMES def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: @@ -2370,7 +2373,15 @@ def refresh_proxy_server_request_body_snapshot( if not isinstance(proxy_server_request, dict): return _body_snapshot_exclude: Final = ( - frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG}) + frozenset( + { + "secret_fields", + "proxy_server_request", + "litellm_logging_obj", + VECTOR_STORE_CALLER_METADATA_KWARG, + STREAM_USAGE_INJECTED_FIELD, + } + ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS | _CALLBACK_CREDENTIAL_KEYS ) @@ -4268,7 +4279,11 @@ _ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join( LlmProviders.VERTEX_AI.value, ) ) +_NATIVE_MESSAGES_ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join( + (LlmProviders.AZURE_AI.value, LlmProviders.DATABRICKS.value) +) _ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = ",".join(sorted(ANTHROPIC_OAUTH_FORWARD_PROVIDERS)) +_CHATGPT_OAUTH_CREDENTIAL_PROVIDERS: Final = LlmProviders.CHATGPT.value def add_provider_specific_headers_to_request( @@ -4276,6 +4291,7 @@ def add_provider_specific_headers_to_request( headers: Mapping[str, str], ) -> bool: from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + from litellm.llms.chatgpt.common_utils import select_chatgpt_client_credential_headers anthropic_api_headers: Final = {header: headers[header] for header in ANTHROPIC_API_HEADERS if header in headers} anthropic_oauth_credential_headers: Final = { @@ -4283,6 +4299,8 @@ def add_provider_specific_headers_to_request( for header, value in headers.items() if header.lower() == "authorization" and is_anthropic_oauth_key(value) } + chatgpt_oauth_credential_headers: Final = dict(select_chatgpt_client_credential_headers(headers)) + typesafe_danger_zone_headers: Final = { header: value for header, value in headers.items() @@ -4296,7 +4314,7 @@ def add_provider_specific_headers_to_request( custom_llm_provider=_ANTHROPIC_API_HEADER_PROVIDERS, extra_headers=anthropic_api_headers ), ProviderSpecificHeader( - custom_llm_provider=LlmProviders.DATABRICKS.value, + custom_llm_provider=_NATIVE_MESSAGES_ANTHROPIC_API_HEADER_PROVIDERS, extra_headers=anthropic_api_headers, native_anthropic_messages_only=True, ), @@ -4307,13 +4325,17 @@ def add_provider_specific_headers_to_request( ProviderSpecificHeader( custom_llm_provider=LlmProviders.TYPESAFE.value, extra_headers=typesafe_danger_zone_headers ), + ProviderSpecificHeader( + custom_llm_provider=_CHATGPT_OAUTH_CREDENTIAL_PROVIDERS, + extra_headers=chatgpt_oauth_credential_headers, + ), ) if scoped_header["extra_headers"] ] if scoped_headers: data["provider_specific_header"] = scoped_headers[0] if len(scoped_headers) == 1 else scoped_headers - return bool(anthropic_oauth_credential_headers) + return bool(anthropic_oauth_credential_headers or chatgpt_oauth_credential_headers) def _add_otel_traceparent_to_data(data: dict, request: Request): diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 5f31aa788e2..a87455e5a60 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -9,6 +9,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager from litellm.proxy._types import ( CommonProxyErrors, @@ -166,6 +167,16 @@ class _AccessGroupTx(Protocol): def litellm_agentstable(self) -> _AgentTable: ... +class _AccessGroupTransaction(Protocol): + async def __aenter__(self) -> _AccessGroupTx: ... + + async def __aexit__(self, *exc_info: object) -> None: ... + + +def _access_group_transaction(prisma_client: PrismaClient) -> _AccessGroupTransaction: + return prisma_client.db.tx() + + def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( @@ -191,6 +202,10 @@ def _tool_permissions_json(value: Mapping[str, Sequence[str]] | None) -> "Json": return Json(None) if value is None else Json({server: list(tools) for server, tools in value.items()}) +def _explicitly_set_fields(data: AccessGroupUpdateRequest) -> Mapping[str, object]: + return data.model_dump(exclude_unset=True) + + @dataclass(frozen=True, slots=True) class _ResourceNames: mcp_servers: Mapping[str, str] @@ -539,7 +554,7 @@ async def create_access_group( try: tx: _AccessGroupTx - async with prisma_client.db.tx() as tx: + async with _access_group_transaction(prisma_client) as tx: existing: Final = await tx.litellm_accessgrouptable.find_unique( where={"access_group_name": data.access_group_name} ) @@ -647,7 +662,7 @@ async def update_access_group( _require_proxy_admin(user_api_key_dict) prisma_client: Final = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - update_fields: Final = data.model_dump(exclude_unset=True) + update_fields: Final = _explicitly_set_fields(data) update_data: Final[dict] = {"updated_by": user_api_key_dict.user_id} for field, value in update_fields.items(): if field in ( @@ -658,7 +673,7 @@ async def update_access_group( "access_agent_ids", "access_skill_names", ): - value = list(value or ()) + value = list(iter_items(value or ())) update_data[field] = ( _tool_permissions_json(data.mcp_tool_permissions) if field == "mcp_tool_permissions" else value ) @@ -672,7 +687,7 @@ async def update_access_group( try: tx: _AccessGroupTx - async with prisma_client.db.tx() as tx: + async with _access_group_transaction(prisma_client) as tx: # Read inside the transaction so delta computation is consistent with the write, # avoiding a TOCTOU race where a concurrent update could make deltas stale. existing: Final = await tx.litellm_accessgrouptable.find_unique(where={"access_group_id": access_group_id}) @@ -687,10 +702,10 @@ async def update_access_group( old_team_ids: Final[set[str]] = set(attached[access_group_id]) old_key_ids: Final[set[str]] = set(existing.assigned_key_ids or []) new_team_ids: Final[set[str]] = ( - set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids + set(data.assigned_team_ids or []) if "assigned_team_ids" in update_fields else old_team_ids ) new_key_ids: Final[set[str]] = ( - set(update_fields["assigned_key_ids"] or []) if "assigned_key_ids" in update_fields else old_key_ids + set(data.assigned_key_ids or []) if "assigned_key_ids" in update_fields else old_key_ids ) teams_to_add = list(new_team_ids - old_team_ids) @@ -747,7 +762,7 @@ async def delete_access_group( affected_key_tokens: list[str] = [] tx: _AccessGroupTx - async with prisma_client.db.tx() as tx: + async with _access_group_transaction(prisma_client) as tx: existing: Final = await tx.litellm_accessgrouptable.find_unique(where={"access_group_id": access_group_id}) if existing is None: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index d40f64f435f..9c46e711cf6 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -54,6 +54,7 @@ from litellm.proxy.management_helpers.auto_router_permissions import ( from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.daily_activity_sql import build_where_clause +from litellm.repositories.table_repositories import PrismaTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository @@ -179,7 +180,7 @@ class _ShadowEvalAttemptRow(Protocol): class _ShadowEvalFunnelTable(Protocol): - async def create_many(self, data: Sequence[Mapping[str, object]], skip_duplicates: bool) -> int: ... + async def create_many(self, data: Sequence[Mapping[str, object]], *, skip_duplicates: bool) -> int: ... class _ShadowEvalAttemptTable(Protocol): @@ -204,16 +205,28 @@ def _user_rows(prisma_client: "PrismaClient") -> _UserRowsTable: return UserRepository(prisma_client).table +class _ShadowEvalJobRepository(PrismaTableRepository[_ShadowEvalJobRow]): + table_name = "litellm_shadowevaljob" + + +class _ShadowEvalFunnelRepository(PrismaTableRepository[object]): + table_name = "litellm_shadowevalfunnel" + + +class _ShadowEvalAttemptRepository(PrismaTableRepository[_ShadowEvalAttemptRow]): + table_name = "litellm_shadowevalattempt" + + def _shadow_eval_jobs(prisma_client: "PrismaClient") -> _ShadowEvalJobTable: - return prisma_client.db.litellm_shadowevaljob + return _ShadowEvalJobRepository(prisma_client).table def _shadow_eval_funnel(prisma_client: "PrismaClient") -> _ShadowEvalFunnelTable: - return prisma_client.db.litellm_shadowevalfunnel # pyright: ignore[reportAttributeAccessIssue] # generated client + return _ShadowEvalFunnelRepository(prisma_client).table def _shadow_eval_attempts(prisma_client: "PrismaClient") -> _ShadowEvalAttemptTable: - return prisma_client.db.litellm_shadowevalattempt + return _ShadowEvalAttemptRepository(prisma_client).table async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]: @@ -221,6 +234,10 @@ async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) - return await prisma_client.db.query_raw(query, *args) +async def _execute_raw(prisma_client: "PrismaClient", query: str, *args: object) -> int: + return await prisma_client.db.execute_raw(query, *args) + + async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> LiteLLM_TeamTable | None: from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, @@ -1788,7 +1805,7 @@ async def start_shadow_eval( # deliberate. Sweep and claim filter on exact (target_type, id) pairs so a team id # that happens to equal a key hash never matches the other kind's slot. for target_type, ids in requested_by_type: - await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, list(ids), target_type) + await _execute_raw(prisma_client, _SWEEP_FINISHED_JOBS_SQL, list(ids), target_type) claimed: Final = await _shadow_eval_jobs(prisma_client).find_many( where={ "OR": [ @@ -2005,8 +2022,8 @@ async def stop_shadow_eval_job( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) stamp: Final = datetime.now(timezone.utc) operator: Final = user_api_key_dict.user_id or "operator" - claimed: Final = await prisma_client.db.execute_raw( - _STOP_JOB_SQL, job_id, operator, stamp.replace(tzinfo=None).isoformat() + claimed: Final = await _execute_raw( + prisma_client, _STOP_JOB_SQL, job_id, operator, stamp.replace(tzinfo=None).isoformat() ) legs: Final = _LEG_ROWS.validate_python( await _shadow_eval_jobs(prisma_client).find_many(where={"group_id": job_id}) or () diff --git a/litellm/proxy/management_endpoints/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py index b1f0002e143..14dd9ce6e8a 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -42,6 +42,7 @@ from litellm.proxy.db.db_span import db_span if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import PrismaClient + from litellm.repositories.prisma_protocols import TableActions from litellm.proxy.common_utils.encrypt_decrypt_utils import ( # noqa: F401 # legacy module exports _ALGO_AES_GCM, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _ENCRYPTION_ALGORITHM_SETTING, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -523,7 +524,7 @@ async def _scan_one_table( scalar_columns: tuple, ) -> LocationReport: report: Final = LocationReport(location=location) - table: Final = getattr(prisma_client.db, db_attr, None) + table: Final[TableActions[object] | None] = getattr(prisma_client.db, db_attr, None) if table is None: return report try: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index b563dbc5c4c..d6bf9220551 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2093,7 +2093,9 @@ async def get_user_key_counts( return result -def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str] | None: +def _validate_sort_params( + sort_by: str | None, sort_order: str +) -> Mapping[str, object] | Sequence[Mapping[str, object]] | None: order_by: Final[dict[str, str]] = {} if sort_by is None: @@ -2106,6 +2108,11 @@ def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str "spend", "user_alias", "user_role", + "max_budget", + "updated_at", + "sso_user_id", + "status", + "key_count", ] if sort_by not in valid_columns: raise HTTPException( @@ -2120,6 +2127,11 @@ def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str detail={"error": "Invalid sort order. Must be 'asc' or 'desc'"}, ) + if sort_by in {"status", "key_count"}: + return None + if sort_by in {"max_budget", "updated_at", "sso_user_id"}: + nulls: Final = "first" if sort_by == "max_budget" and sort_order.lower() == "desc" else "last" + return ({sort_by: {"sort": sort_order.lower(), "nulls": nulls}}, {"user_id": "asc"}) order_by[sort_by] = sort_order.lower() return order_by @@ -2338,16 +2350,40 @@ async def get_users( # Build order_by conditions - order_by: Final[dict[str, str] | None] = ( + order_by: Final = ( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) - users: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await UserRepository(prisma_client).table.find_many( - where=where, - skip=skip, - take=page_size, - order=(order_by if order_by else {"created_at": "desc"}), # Default to created_at desc if no sort specified + from litellm.proxy.management_helpers.user_list_sorting import sorted_user_ids + + computed_sort: Final = sort_by in ("status", "key_count") + sorted_ids: Final = ( + await sorted_user_ids( + prisma_client, + sort_by="status" if sort_by == "status" else "key_count", + descending=sort_order.lower() == "desc", + skip=skip, + take=page_size, + role=role, + user_ids=user_ids, + user_email=user_email, + team=team, + sso_user_ids=sso_user_ids, + organization_ids=organization_ids, + search=search, + ) + if computed_sort + else () ) + page_users: Final = await UserRepository(prisma_client).table.find_many( + # bounded-ok: sorted_ids comes from LIMIT page_size, capped at 100 by the route + where={"AND": (where, {"user_id": {"in": sorted_ids}})} if computed_sort else where, + skip=0 if computed_sort else skip, + take=page_size, + order=order_by if order_by else {"created_at": "desc"}, + ) + positions: Final = {user_id: index for index, user_id in enumerate(sorted_ids)} + users: Final = sorted(page_users, key=lambda user: positions[user.user_id]) if computed_sort else page_users # Get total count of user rows total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 824b310a1dd..0ad4258fb04 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -211,6 +211,7 @@ if TYPE_CHECKING: from prisma import models as prisma_models _RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel) +_RawRowT = TypeVar("_RawRowT") class _UserRowLike(Protocol): @@ -234,6 +235,12 @@ class _TxTables(Protocol): litellm_proxymodeltable: TableActions[object] +class _ModelReencryptTransaction(Protocol): + async def __aenter__(self) -> _TxTables: ... + + async def __aexit__(self, *exc_info: object) -> None: ... + + class _ModelParamsUpdate(TypedDict): litellm_params: ReadOnly["prisma.Json"] @@ -334,6 +341,14 @@ def _credentials_table(prisma_client: PrismaClient) -> TableActions[CredentialIt ) +async def _query_raw_rows(prisma_client: PrismaClient, query: str, *args: object) -> Sequence[_RawRowT]: + return await prisma_client.db.query_raw(query, *args) + + +def _model_reencrypt_transaction(prisma_client: PrismaClient) -> _ModelReencryptTransaction: + return prisma_client.db.tx(timeout=timedelta(minutes=2)) + + 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, use_writer=use_writer).table @@ -1313,11 +1328,6 @@ async def validate_team_id_used_in_service_account_request( return True -_BUDGET_NUMERIC_KEYS = frozenset( - ["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit", "tpd_limit"] -) - - def _enforce_upperbound_key_params( data: GenerateKeyRequest | UpdateKeyRequest, fill_defaults: bool = True, @@ -1332,13 +1342,20 @@ def _enforce_upperbound_key_params( # (GHSA-2rv4-xv66-fpjg): float('nan') passes every `< 0` check because # nan < 0 is False, and spend >= nan is always False, permanently disabling # budget enforcement for any key that carries it. - for elem in data: - key, value = elem - if key in _BUDGET_NUMERIC_KEYS and value is not None: - if not math.isfinite(value): + budget_numeric_fields: Final = ( + ("max_budget", data.max_budget), + ("max_parallel_requests", data.max_parallel_requests), + ("tpm_limit", data.tpm_limit), + ("rpm_limit", data.rpm_limit), + ("tpd_limit", data.tpd_limit), + ("soft_budget", data.soft_budget), + ) + for field_name, field_value in budget_numeric_fields: + if field_value is not None: + if not math.isfinite(field_value): raise HTTPException( status_code=400, - detail={"error": f"{key} must be a finite number. Received: {value}"}, + detail={"error": f"{field_name} must be a finite number. Received: {field_value}"}, ) if litellm.upperbound_key_generate_params is None: @@ -5746,8 +5763,7 @@ async def rotate_master_key( ] ) verbose_proxy_logger.debug("Re-encrypting litellm_params on %s model rows", len(reencrypted_models)) - async with prisma_client.db.tx(timeout=timedelta(minutes=2)) as tx_ctx: - tx: Final[_TxTables] = tx_ctx + async with _model_reencrypt_transaction(prisma_client) as tx: for reencrypted_model in reencrypted_models: await tx.litellm_proxymodeltable.update_many( data=_ModelParamsUpdate(litellm_params=prisma.Json(reencrypted_model.litellm_params)), @@ -7344,7 +7360,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final[Sequence[Mapping[str, int]]] = await _query_raw_rows(prisma_client, count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -7357,7 +7373,9 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final[Sequence[Mapping[str, str]]] = await _query_raw_rows( + prisma_client, aliases_sql, *aliases_params + ) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 @@ -7397,20 +7415,26 @@ async def key_aliases( ) -def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str] | None: - order_by: Final[dict[str, str]] = {} - +def _validate_sort_params( + sort_by: str | None, sort_order: str +) -> Mapping[str, object] | Sequence[Mapping[str, object]] | None: if sort_by is None: return None # Validate sort_by is a valid column - valid_columns: Final = [ + valid_columns: Final = ( "spend", "max_budget", "created_at", "updated_at", "token", "key_alias", - ] + "expires", + "last_active", + "budget_reset_at", + "total_spend", + "tpm_limit", + "rpm_limit", + ) if sort_by not in valid_columns: raise HTTPException( status_code=400, @@ -7424,9 +7448,14 @@ def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str detail={"error": "Invalid sort order. Must be 'asc' or 'desc'"}, ) - order_by[sort_by] = sort_order.lower() - - return order_by + direction: Final = sort_order.lower() + nullable_sort_columns: Final = ("expires", "last_active", "budget_reset_at", "tpm_limit", "rpm_limit") + if sort_by not in (*nullable_sort_columns, "total_spend"): + return {sort_by: direction} + return [ + {sort_by: {"sort": direction, "nulls": "last"}} if sort_by in nullable_sort_columns else {sort_by: direction}, + {"token": "asc"}, + ] def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, object]: @@ -7668,7 +7697,7 @@ async def list_key_helper( verbose_proxy_logger.debug("Pagination: skip=%s, take=%s", skip, size) - order_by: Final[dict[str, str] | None] = ( + order_by: Final = ( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py index 1970a0797ef..19a5408496a 100644 --- a/litellm/proxy/management_endpoints/management_v1/budgets.py +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -3,6 +3,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime +from itertools import chain from types import MappingProxyType from typing import Annotated, Final @@ -76,6 +77,32 @@ _ROW_COUNTS: Final = TypeAdapter(tuple[_RowCount, ...]) SELECTED_COLUMNS: Final = ", ".join(f'"{name}"' for name in BudgetListItem.model_fields) +# Unicode 17.0 Nd zeroes, each followed by its other nine digits in UnicodeData.txt. +_DECIMAL_ZEROES: Final = "0٠۰߀०০੦૦୦௦౦೦൦෦๐໐༠၀႐០᠐᥆᧐᪀᪐᭐᮰᱀᱐꘠꣐꤀꧐꧰꩐꯰0𐒠𐴰𐵀𑁦𑃰𑄶𑇐𑋰𑑐𑓐𑙐𑛀𑛐𑛚𑜰𑣠𑥐𑯰𑱐𑵐𑶠\U00011de0𑽐𖄰𖩠𖫀𖭐𖵰𜳰𝟎𝟘𝟢𝟬𝟶𞅀𞋰𞓰𞗱𞥐🯰" +_DECIMAL_DIGITS: Final = "".join( + chain.from_iterable(tuple(chr(ord(zero) + digit) for digit in range(10)) for zero in _DECIMAL_ZEROES) +) +_ASCII_DIGITS: Final = "0123456789" * len(_DECIMAL_ZEROES) + + +_BUDGET_DURATION_SECONDS: Final = f"""( + SELECT EXTRACT(EPOCH FROM ( + (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') + + (parts[1] || ' ' || CASE parts[2] + WHEN 's' THEN 'seconds' WHEN 'm' THEN 'minutes' WHEN 'h' THEN 'hours' + WHEN 'd' THEN 'days' WHEN 'w' THEN 'weeks' WHEN 'mo' THEN 'months' END)::interval + - (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') + )) FROM regexp_match( + CASE lower(trim(budget_duration)) + WHEN 'hourly' THEN '1h' WHEN 'daily' THEN '24h' + WHEN 'weekly' THEN '7d' WHEN 'monthly' THEN '30d' + ELSE CASE WHEN octet_length(budget_duration) = length(budget_duration) THEN budget_duration + ELSE translate(budget_duration, '{_DECIMAL_DIGITS}', '{_ASCII_DIGITS}') END END, + '^([0-9]+)(mo|[smhdw])' + ) AS duration(parts) +)""" + + @dataclass(frozen=True, slots=True) class PrismaBudgetListExecutor: """The database half of the budgets list. Every caller-supplied value is bound to a @@ -91,10 +118,16 @@ class PrismaBudgetListExecutor: async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]: clauses, params = where_sql(plan.where) + ordering: Final = ", ".join( + order_by_sql((key,)).replace("NULLS LAST", "NULLS FIRST") + if key.descending and key.field in ("max_budget", "tpm_limit", "rpm_limit", "tpd_limit") + else order_by_sql((key,)) + for key in plan.order + ).replace('"budget_duration"', _BUDGET_DURATION_SECONDS) sql: Final = ( f"SELECT {SELECTED_COLUMNS} FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "") - + f" ORDER BY {order_by_sql(plan.order)}" + + f" ORDER BY {ordering}" + f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" ) return _BUDGET_ROWS.validate_python(await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip)) @@ -111,8 +144,6 @@ def _scope(caller: UserAPIKeyAuth) -> Scope: return ScopeDenied(reason=f"Only proxy admins can list budgets, your role={caller.user_role}") -# budget_duration is deliberately absent from `sortable`: the column holds strings -# like "7d" and "30d", so a lexicographic ORDER BY puts "30d" ahead of "7d". BUDGET_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType( { "budget_duration": FilterSpec(type=str, ops=frozenset(("in", "is_null"))), @@ -123,7 +154,9 @@ BUDGET_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType( BUDGETS_LIST_SPEC: Final[ListSpec[BudgetListItem, BudgetListItem]] = ListSpec( resource="budgets", - sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at")), + sortable=frozenset( + ("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at", "budget_duration") + ), searchable=frozenset(("budget_id",)), filters=BUDGET_FILTERS, default_sort=(SortKey(field="created_at", descending=True),), @@ -154,8 +187,9 @@ async def list_budgets( way to page, sort or filter it. `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, - `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending, - and defaults to `-created_at`. `budget_id` is appended to every sort as the + `rpm_limit`, `tpd_limit`, `created_at` or `budget_duration`, each optionally prefixed with `-` for descending, + and defaults to `-created_at`. Reset periods sort by elapsed time, with calendar months + measured from the current UTC date and unset periods last. `budget_id` is appended to every sort as the tiebreaker. `q` is a case-insensitive substring match on `budget_id`. `page_size` defaults to 50 and is capped at 100. Filters are `filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]` diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index b93c05c555f..da84c2d0013 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -1,10 +1,12 @@ """`/management/v1/spend_logs` facets.""" +from collections.abc import Sequence from datetime import datetime, timezone from functools import partial -from typing import Annotated, Final, Literal +from typing import Annotated, Final, Literal, Protocol from fastapi import APIRouter, Depends, Query, Request +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth @@ -19,6 +21,7 @@ from litellm.proxy.list_api.common import ( reject_unknown_query_params, ) from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX +from litellm.proxy.utils import PrismaClient from litellm.types.proxy.management_endpoints.management_v1 import ( FacetListResponse, PageMeta, @@ -33,6 +36,19 @@ router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX) SPEND_LOGS_FACET_SCAN_CAP: Final = 10000 +class _FacetRow(TypedDict, total=False): + end_user: ReadOnly[str | None] + user: ReadOnly[str | None] + + +class _FacetQuery(Protocol): + async def __call__(self, query: str, *args: object) -> Sequence[_FacetRow]: ... + + +def _facet_query(prisma_client: PrismaClient) -> _FacetQuery: + return prisma_client.db.query_raw + + def _as_utc(value: datetime) -> datetime: return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) @@ -125,8 +141,8 @@ async def _list_spend_log_facet( f" ORDER BY {column_sql} ASC" f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}" ) - rows: Final = await prisma_client.db.query_raw(facet_sql, *params) - values: Final[list[str]] = [row[column] for row in rows if row.get(column)] + rows: Final = await _facet_query(prisma_client)(facet_sql, *params) + values: Final[list[str]] = [value for row in rows if (value := row.get(column))] has_more: Final = len(values) > page_size return FacetListResponse( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a014b9cf8a5..776072c69be 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -104,6 +104,14 @@ class _HasServerId(Protocol): server_id: str +class _TemporaryServerCacheWriter(Protocol): + async def async_set_cache(self, *, key: str, value: str, ttl: int) -> object: ... + + +class _TemporaryServerCacheReader(Protocol): + async def async_get_cache(self, *, key: str) -> object: ... + + def _session_cookie_claims(claims: Mapping[str, object]) -> Mapping[str, object]: return claims @@ -527,7 +535,7 @@ if MCP_AVAILABLE: """ if litellm.cache is None or not hasattr(litellm.cache, "cache"): return - cache_backend: Final = getattr(litellm.cache, "cache", None) + cache_backend: Final[_TemporaryServerCacheWriter | None] = getattr(litellm.cache, "cache", None) if cache_backend is None or not hasattr(cache_backend, "async_set_cache"): return @@ -566,7 +574,7 @@ if MCP_AVAILABLE: """ if litellm.cache is None or not hasattr(litellm.cache, "cache"): return None - cache_backend: Final = getattr(litellm.cache, "cache", None) + cache_backend: Final[_TemporaryServerCacheReader | None] = getattr(litellm.cache, "cache", None) if cache_backend is None or not hasattr(cache_backend, "async_get_cache"): return None diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index fc0868a5fc9..11f616a5e18 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -3217,12 +3217,13 @@ async def update_useful_links( ) -def _validated_labeled_tiers( - tier_labels: dict[ComplexityTier, str], # mutable-ok: Pydantic materializes JSON object fields as dicts -) -> tuple[tuple[ComplexityTier, str], ...]: +_DECODED_JSON: Final = TypeAdapter(object) + + +def _validated_labeled_tiers(tier_labels: object) -> tuple[tuple[ComplexityTier, str], ...]: """Validate tier labels once for both prompt-preview transports.""" try: - return ComplexityRouterConfig(tier_labels=tier_labels).labeled_tiers() + return ComplexityRouterConfig.model_validate({"tier_labels": tier_labels}).labeled_tiers() except (TypeError, ValidationError) as e: raise ProxyException( message=f"tier_labels must be a JSON object of tier name to display name: {e}", @@ -3237,7 +3238,7 @@ def _labeled_tiers_from_query(tier_labels: str | None) -> tuple[tuple[Complexity if not tier_labels: return None try: - parsed: Final = json.loads(tier_labels) + parsed: Final = _DECODED_JSON.validate_python(json.loads(tier_labels)) except JSONDecodeError as e: raise ProxyException( message=f"tier_labels must be a JSON object of tier name to display name: {e}", diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 3bef7c107eb..7cd36d51511 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -28,6 +28,7 @@ from litellm.constants import ( MAX_COMPETITOR_NAMES, ) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) @@ -657,7 +658,7 @@ async def get_policy_templates( Fetches from GitHub with automatic fallback to local backup on failure. Set LITELLM_LOCAL_POLICY_TEMPLATES=true to skip GitHub and use local backup only. """ - use_local: Final = os.getenv("LITELLM_LOCAL_POLICY_TEMPLATES", "").strip().lower() in ( + use_local: Final = is_offline_mode() or os.getenv("LITELLM_LOCAL_POLICY_TEMPLATES", "").strip().lower() in ( "true", "1", "yes", diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 64566221d4d..1f4ca442d20 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -71,10 +71,11 @@ def _get_routing_strategies_from_router_class() -> list[str]: # Get the routing_strategy parameter routing_strategy_param: Final = sig.parameters.get("routing_strategy") + routing_strategy_annotation: Final[object] = routing_strategy_param.annotation if routing_strategy_param else None - if routing_strategy_param and routing_strategy_param.annotation: + if routing_strategy_param and routing_strategy_annotation: # Extract Literal values using get_args - literal_values: Final = get_args(routing_strategy_param.annotation) + literal_values: Final = get_args(routing_strategy_annotation) if literal_values: return list(literal_values) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d9763f72741..1fa18e9b2b8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -52,6 +52,7 @@ from litellm.proxy._types import ( CommonProxyErrors, DeleteTeamRequest, LiteLLM_AuditLogs, + LiteLLM_BudgetTable, LiteLLM_DeletedTeamTable, Litellm_EntityType, LiteLLM_ManagementEndpoint_MetadataFields, @@ -114,6 +115,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo get_user_object, invalidate_team_member_spend_state, jwt_key_mapping_cache_key, + resolve_effective_team_member_budget, ) from litellm.proxy.auth.auth_utils import ( enforce_batch_limits_are_admin_only, @@ -244,6 +246,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, TeamIdSearchMatch, + TeamListCallerMembership, TeamListItem, TeamListResponse, TeamMemberAddResult, @@ -254,6 +257,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( UpdateTeamMemberPermissionsRequest, ) from litellm.types.utils import BudgetConfig, GenericBudgetConfigType +from litellm.utils import get_utc_datetime if TYPE_CHECKING: from prisma import Prisma @@ -302,6 +306,34 @@ class _TeamIdRow(Protocol): def team_id(self) -> str: ... +class _TeamListBudgetRow(Protocol): + @property + def budget_id(self) -> str: ... + + @property + def budget_reset_at(self) -> datetime | None: ... + + @property + def max_budget(self) -> float | None: ... + + @property + def temp_budget_expiry(self) -> datetime | None: ... + + @property + def temp_budget_increase(self) -> float | None: ... + + +class _TeamListMembershipRow(Protocol): + @property + def team_id(self) -> str: ... + + @property + def spend(self) -> float | None: ... + + @property + def litellm_budget_table(self) -> _TeamListBudgetRow | None: ... + + class _CacheableTeamRow(_TeamIdRow, _ModelDumpRow, Protocol): ... @@ -479,6 +511,14 @@ def _team_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models. return TeamRepository(prisma_client).table +class _ExecuteRaw(Protocol): + async def __call__(self, query: str, *args: object) -> int: ... + + +def _execute_raw(prisma_client: PrismaClient) -> _ExecuteRaw: + return prisma_client.db.execute_raw + + def _team_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamTable]": return cast( # cast-ok: generated actions type Json columns as str; TableActions widens inputs to Mapping "TableActions[prisma_models.LiteLLM_TeamTable]", tx.litellm_teamtable @@ -973,7 +1013,7 @@ async def get_all_team_memberships( returned_tm: Final[list[LiteLLM_TeamMembership]] = [] for tm in team_memberships: - returned_tm.append(LiteLLM_TeamMembership.model_validate(tm.model_dump())) + returned_tm.append(with_total_spend_floor(LiteLLM_TeamMembership.model_validate(tm.model_dump()))) return returned_tm @@ -3092,7 +3132,7 @@ async def _process_team_members( # Resolve allowed_models: explicit request value, or fall back to team's default_team_member_models member_allowed_models = data.allowed_models - team_default_member_models: Final = getattr(complete_team_data, "default_team_member_models", None) + team_default_member_models: Final = complete_team_data.default_team_member_models if member_allowed_models is None and team_default_member_models: member_allowed_models = team_default_member_models @@ -4162,6 +4202,23 @@ def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAu _raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.") +# One statement so the validated spend cannot go stale between the read and the write: the WHERE re-checks +# reset_to <= spend, and total_spend is floored at spend before spend drops so the history survives the reset. +_RESET_TEAM_MEMBER_SPEND_SQL: Final = """ +UPDATE "LiteLLM_TeamMembership" +SET total_spend = GREATEST(total_spend, spend), spend = $3::float8 +WHERE user_id = $1 AND team_id = $2 AND spend >= $3::float8 +""" + + +def with_total_spend_floor(membership: LiteLLM_TeamMembership) -> LiteLLM_TeamMembership: + """Rows written before total_spend existed (migration 20260421135425) hold their history only in spend.""" + spend: Final = membership.spend or 0.0 + if (membership.total_spend or 0.0) >= spend: + return membership + return membership.model_copy(update={"total_spend": spend}) + + def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn: detail: Final = {"error": message} raise HTTPException(status_code=status_code, detail=detail) @@ -4246,10 +4303,12 @@ async def reset_team_member_spend_fn( current_spend: Final = membership.spend or 0.0 reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership) - await _team_membership_db(prisma_client).update( - where=membership_where, - data={"spend": reset_to}, - ) + reset_rows: Final = await _execute_raw(prisma_client)(_RESET_TEAM_MEMBER_SPEND_SQL, user_id, team_id, reset_to) + if reset_rows == 0: + _raise_reset_spend_error( + status.HTTP_409_CONFLICT, + f"Member spend changed since it was read; reset_to ({reset_to}) must be <= current spend. Retry.", + ) await invalidate_team_member_spend_state( user_id=user_id, @@ -5459,16 +5518,15 @@ async def team_member_me( user_api_key_cache=user_api_key_cache, ) - caller_user_email: Final = user_api_key_dict.user_email - member_role: str | None = None - for m in team_table.members_with_roles: - # Match by user_id when present, else fall back to email — members - # added by email may have user_id=None on the stored entry. - if (m.user_id is not None and m.user_id == caller_user_id) or ( - m.user_email is not None and caller_user_email is not None and m.user_email == caller_user_email - ): - member_role = m.role - break + caller_identity: Final = _TeamListCallerIdentity( + user_id=caller_user_id, + user_email=user_api_key_dict.user_email, + ) + matched_member: Final = next( + (member for member in team_table.members_with_roles if _team_member_matches_caller(member, caller_identity)), + None, + ) + member_role: Final = matched_member.role if matched_member is not None else None if member_role is None: # Caller is not a member of this team. Even proxy admins get 404 here — @@ -5509,16 +5567,17 @@ async def team_member_me( litellm_budget_table=None, ) + floored: Final = with_total_spend_floor(membership) return TeamMemberInfoResponse( user_id=caller_user_id, team_id=team_id, team_alias=team_table.team_alias, role=member_role, user_email=user_email, - spend=membership.spend, - total_spend=membership.total_spend, - budget_id=membership.budget_id, - litellm_budget_table=membership.litellm_budget_table, + spend=floored.spend, + total_spend=floored.total_spend, + budget_id=floored.budget_id, + litellm_budget_table=floored.litellm_budget_table, ) @@ -5878,11 +5937,16 @@ def _convert_teams_to_response_models( members_count = len(members_with_roles) keys_count = counts.get(team_dict.get("team_id") or "", 0) team_list.append( - TeamListItem( - **team_dict, - members_count=members_count, - keys_count=keys_count, - access_group_skill_names=None, + TeamListItem.model_validate( + cast( # cast-ok: TeamListItem validates the Prisma model dump + object, + { + **team_dict, + "members_count": members_count, + "keys_count": keys_count, + "access_group_skill_names": None, + }, + ) ) ) return team_list @@ -5913,6 +5977,137 @@ async def _get_keys_count_by_team( return {row["team_id"]: row.get("_count", {}).get("team_id", 0) for row in grouped if row.get("team_id")} +def _team_list_budget_model(budget_row: _TeamListBudgetRow | None) -> LiteLLM_BudgetTable | None: + if budget_row is None: + return None + return LiteLLM_BudgetTable( + max_budget=budget_row.max_budget, + temp_budget_expiry=budget_row.temp_budget_expiry, + temp_budget_increase=budget_row.temp_budget_increase, + ) + + +@dataclass(frozen=True, slots=True) +class _TeamListCallerIdentity: + user_id: str | None + user_email: str | None + + +def _team_member_matches_caller(member: Member, caller_identity: _TeamListCallerIdentity) -> bool: + return (member.user_id is not None and member.user_id == caller_identity.user_id) or ( + member.user_email is not None + and caller_identity.user_email is not None + and member.user_email == caller_identity.user_email + ) + + +def _caller_in_roster(team: TeamListItem, caller_identity: _TeamListCallerIdentity) -> bool: + return any(_team_member_matches_caller(member, caller_identity) for member in team.members_with_roles) + + +def _team_member_default_budget_id( + team: TeamListItem, + caller_identity: _TeamListCallerIdentity, + memberships_by_team: Mapping[str, _TeamListMembershipRow], +) -> str | None: + membership_row: Final = memberships_by_team.get(team.team_id) + if membership_row is None and not _caller_in_roster(team, caller_identity): + return None + member_budget: Final = _team_list_budget_model( + membership_row.litellm_budget_table if membership_row is not None else None + ) + if member_budget is not None and member_budget.max_budget is not None: + return None + metadata: Final = cast( # cast-ok: team metadata is stored as a JSON object + "dict[str, object] | None", team.metadata + ) + if metadata is None: + return None + budget_id: Final = metadata.get("team_member_budget_id") + return budget_id if isinstance(budget_id, str) else None + + +def _team_list_caller_membership( + team: TeamListItem, + caller_identity: _TeamListCallerIdentity, + memberships_by_team: Mapping[str, _TeamListMembershipRow], + default_budget_id: str | None, + budgets_by_id: Mapping[str, _TeamListBudgetRow], + now: datetime, +) -> TeamListCallerMembership | None: + membership_row: Final = memberships_by_team.get(team.team_id) + if membership_row is None and not _caller_in_roster(team, caller_identity): + return None + member_budget_row: Final = membership_row.litellm_budget_table if membership_row is not None else None + member_budget: Final = _team_list_budget_model(member_budget_row) + default_budget_row: Final = budgets_by_id.get(default_budget_id) if default_budget_id is not None else None + team_default_budget: Final = _team_list_budget_model(default_budget_row) + max_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget, + team_default_budget=team_default_budget, + now=now, + ) + budget_reset_at: Final = ( + member_budget_row.budget_reset_at if member_budget_row is not None and max_budget is not None else None + ) + return TeamListCallerMembership( + spend=(membership_row.spend if membership_row is not None else 0.0) or 0.0, + max_budget=max_budget, + budget_reset_at=budget_reset_at, + ) + + +async def _enrich_team_list_with_caller_membership( + prisma_client: PrismaClient, + teams: Sequence[TeamListItem], + caller_identity: _TeamListCallerIdentity, + now: datetime, +) -> tuple[TeamListItem, ...]: + caller_user_id: Final = caller_identity.user_id + if caller_user_id is None: + return tuple(team.model_copy(update={"caller_membership": None}) for team in teams) + if not teams: + return () + + page_team_ids: Final = tuple(team.team_id for team in teams) + membership_rows: Final[Sequence[_TeamListMembershipRow]] = await _team_membership_db(prisma_client).find_many( + # bounded-ok: page_team_ids is capped at 100 teams by list_team_v2 + where={"team_id": {"in": page_team_ids}, "user_id": caller_user_id}, + include={"litellm_budget_table": True}, + ) + memberships_by_team: Final = {row.team_id: row for row in membership_rows} + default_budget_id_by_team: Final = MappingProxyType( + {team.team_id: _team_member_default_budget_id(team, caller_identity, memberships_by_team) for team in teams} + ) + default_budget_ids: Final = tuple( + dict.fromkeys(budget_id for budget_id in default_budget_id_by_team.values() if budget_id is not None) + ) + default_budget_rows: Final[Sequence[_TeamListBudgetRow]] = ( + await _budget_db(prisma_client).find_many( + # bounded-ok: default_budget_ids is derived from this page of at most 100 teams + where={"budget_id": {"in": default_budget_ids}}, + ) + if default_budget_ids + else () + ) + budgets_by_id: Final = {row.budget_id: row for row in default_budget_rows} + return tuple( + team.model_copy( + update={ + "caller_membership": _team_list_caller_membership( + team=team, + caller_identity=caller_identity, + memberships_by_team=memberships_by_team, + default_budget_id=default_budget_id_by_team.get(team.team_id), + budgets_by_id=budgets_by_id, + now=now, + ) + } + ) + for team in teams + ) + + async def _enforce_list_team_v2_access( user_api_key_dict: UserAPIKeyAuth, user_id: str | None, @@ -6026,7 +6221,7 @@ async def list_team_v2( sort_order: str = fastapi.Query(default="asc", description="Sort order ('asc' or 'desc')"), status: str | None = fastapi.Query(default=None, description="Filter by status (e.g. 'deleted')"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> TeamListResponse: """ Get a paginated list of teams with filtering and sorting options. @@ -6101,21 +6296,60 @@ async def list_team_v2( ) if where_conditions is None: - return { - "teams": [], - "total": 0, - "page": page, - "page_size": page_size, - "total_pages": 0, - } + return TeamListResponse( + teams=[], + total=0, + page=page, + page_size=page_size, + total_pages=0, + ) - # Build order_by conditions - valid_sort_columns: Final = ["team_id", "team_alias", "created_at"] - order_by = None - if sort_by and sort_by in valid_sort_columns: - if sort_order.lower() not in ["asc", "desc"]: - sort_order = "asc" - order_by = {sort_by: sort_order.lower()} + from litellm.proxy.management_helpers.team_list_sorting import ( + COMPUTED_TEAM_SORTS, + UnsupportedTeamFilter, + sorted_team_ids, + ) + + valid_sort_columns: Final = ( + "team_id", + "team_alias", + "created_at", + "spend", + "max_budget", + "tpm_limit", + "rpm_limit", + "updated_at", + ) + sort_direction: Final = "desc" if str(sort_order).lower() == "desc" else "asc" + nullable_sort: Final = sort_by in ("max_budget", "tpm_limit", "rpm_limit") + order_by: Final = ( + ( + {sort_by: {"sort": sort_direction, "nulls": "first" if sort_direction == "desc" else "last"}}, + {"team_id": "asc"}, + ) + if nullable_sort + else ({sort_by: sort_direction}, {"team_id": "asc"}) + if sort_by in ("spend", "updated_at") + else {sort_by: sort_direction} + if sort_by in valid_sort_columns + else None + ) + computed_sort: Final = not use_deleted_table and sort_by in COMPUTED_TEAM_SORTS + sorted_ids: Final = ( + await sorted_team_ids( + prisma_client, + where=where_conditions, + sort_by=sort_by or "", + descending=sort_direction == "desc", + skip=skip, + take=page_size, + ) + if computed_sort + else () + ) + + if isinstance(sorted_ids, UnsupportedTeamFilter): + raise HTTPException(status_code=500, detail="Unsupported team sorting filter") # Get teams with pagination if use_deleted_table: @@ -6130,26 +6364,32 @@ async def list_team_v2( total_count = await _deleted_team_db(prisma_client).count(where=where_conditions) else: teams = await _team_db(prisma_client).find_many( - where=where_conditions, - skip=skip, + # bounded-ok: sorted_ids comes from LIMIT page_size, capped at 100 by the route + where={"AND": (where_conditions, {"team_id": {"in": sorted_ids}})} if computed_sort else where_conditions, + skip=0 if computed_sort else skip, take=page_size, - order=order_by if order_by else {"created_at": "desc"}, # Default sort + order=order_by if order_by else {"created_at": "desc"}, include=_INCLUDE_MODEL_TABLE, ) # Get total count for pagination total_count = await _team_db(prisma_client).count(where=where_conditions) + positions: Final = {team_id: index for index, team_id in enumerate(sorted_ids)} + sorted_teams: Final = sorted(teams, key=lambda team: positions[team.team_id]) if computed_sort else teams + # Calculate total pages total_pages: Final = -(-total_count // page_size) # Ceiling division # Aggregate virtual-key counts per team for the current page. The deleted # table does not carry keys_count, so it is skipped. - keys_count_by_team: dict[str, int] = {} - if not use_deleted_table: - keys_count_by_team = await _get_keys_count_by_team(prisma_client, teams) + keys_count_by_team: Final = ( + await _get_keys_count_by_team(prisma_client, sorted_teams) if not use_deleted_table else {} + ) # Convert Prisma models to response models with members_count and keys_count - team_list = _convert_teams_to_response_models(teams, use_deleted_table, keys_count_by_team=keys_count_by_team) + team_list: Final = _convert_teams_to_response_models( + sorted_teams, use_deleted_table, keys_count_by_team=keys_count_by_team + ) # Resolve resources inherited from access groups (single batch query) if not use_deleted_table: @@ -6174,13 +6414,27 @@ async def list_team_v2( team_groups, tuple(ag_id for ag_id in (team_item.access_group_ids or []) if ag_id not in ag_lookup) ) - return { - "teams": team_list, - "total": total_count, - "page": page, - "page_size": page_size, - "total_pages": total_pages, - } + caller_identity: Final = _TeamListCallerIdentity( + user_id=user_api_key_dict.user_id, + user_email=user_api_key_dict.user_email, + ) + team_list_with_caller_membership: Final = ( + await _enrich_team_list_with_caller_membership( + prisma_client=prisma_client, + teams=tuple(team for team in team_list if isinstance(team, TeamListItem)), + caller_identity=caller_identity, + now=get_utc_datetime(), + ) + if not use_deleted_table and team_list + else team_list + ) + return TeamListResponse( + teams=list(team_list_with_caller_membership), + total=total_count, + page=page, + page_size=page_size, + total_pages=total_pages, + ) async def _authorize_and_filter_teams( @@ -6527,7 +6781,7 @@ async def append_team_models( # When the team currently has models=[] (unrestricted access), the # CASE expression inserts the 'all-proxy-models' sentinel first. models_to_add: Final = list(data.models) - await prisma_client.db.execute_raw( + await _execute_raw(prisma_client)( 'UPDATE "LiteLLM_TeamTable" ' "SET models = (" " SELECT ARRAY(SELECT DISTINCT unnest(" diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 4414eed97b2..221a6221fab 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -5,7 +5,7 @@ Might include fastapi/proxy requirements.txt related imports """ from collections.abc import Iterable, Sequence -from typing import Any, Final, cast +from typing import Final, cast from fastapi_sso.sso.base import OpenID @@ -94,4 +94,4 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: class CustomOpenID(OpenID): team_ids: list[str] user_role: LitellmUserRoles | None = None - extra_fields: dict[str, Any] | None = None + extra_fields: dict[str, object] | None = None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 19bca3d2bf8..02494de8c87 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -152,7 +152,7 @@ from litellm.proxy.utils import ( from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SSOConfigRepository from litellm.repositories.team_repository import TeamRepository -from litellm.repositories.user_repository import UserRepository +from litellm.repositories.user_repository import UserRepository, free_sso_user_creation from litellm.secret_managers.main import get_secret_bool, get_secret_str from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.management_endpoints.ui_sso import * # noqa: F403 @@ -278,6 +278,12 @@ class _CustomSsoCall(Protocol): async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... +def _configured_custom_sso() -> _CustomSsoCall | None: + from litellm.proxy.proxy_server import user_custom_sso + + return user_custom_sso + + class _ServicePrincipalAssignment(Protocol): def get(self, key: str) -> str: ... @@ -1051,6 +1057,11 @@ class _UserinfoFallbackCapableSSO(Protocol): ) -> CustomOpenID | None: ... +class _GenericSSOClient(_UserinfoFallbackCapableSSO, Protocol): + @property + def refresh_token(self) -> str | None: ... + + async def _verify_generic_sso_or_fall_back_to_id_token( generic_sso: _UserinfoFallbackCapableSSO, request: Request, @@ -1074,25 +1085,8 @@ async def _verify_generic_sso_or_fall_back_to_id_token( ) -async def _raise_if_sso_exceeds_free_user_limit(premium_user: bool, prisma_client: PrismaClient | None) -> None: - """Free tier allows SSO for up to 5 billable users; beyond that requires an Enterprise license.""" - if premium_user is True: - return - if prisma_client is None: - raise ProxyException( - message=CommonProxyErrors.db_not_connected_error.value, - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) - billable_users: Final = await UserRepository(prisma_client).count_billable_users() - if billable_users and billable_users > 5: - raise ProxyException( - message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You configured SSO (one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, `GENERIC_CLIENT_ID`, or SAML) in your env. Please unset it", - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) +class _SSOUserAdmissionError(ProxyException): + pass @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) @@ -1112,8 +1106,6 @@ async def google_login( from litellm.proxy.proxy_server import ( cli_sso_session_cache, general_settings, - premium_user, - prisma_client, saml_sso_cache, user_custom_ui_sso_sign_in_handler, ) @@ -1128,15 +1120,6 @@ async def google_login( if admin_ui_is_disabled and not is_cli_sso_login: return admin_ui_disabled() - ####### Check if user is a Enterprise / Premium User ####### - if ( - microsoft_client_id is not None - or google_client_id is not None - or generic_client_id is not None - or SAMLAuthHandler.is_saml_configured() - ): - await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) - ####### Detect DB + MASTER KEY in .env ####### missing_env_vars: Final = show_missing_vars_in_env() if missing_env_vars is not None: @@ -1606,7 +1589,7 @@ async def get_generic_sso_response( team_mappings: Final = await _setup_team_mappings() generic_include_token_claims: Final = os.getenv("GENERIC_INCLUDE_TOKEN_CLAIMS", "false").lower() == "true" - def response_convertor(response: Mapping[str, object], httpx_session: object): + def response_convertor(response: Mapping[str, object], httpx_session: object) -> CustomOpenID: nonlocal received_response # return for user debugging response_id_token: Final = response.get("id_token") response_access_token: Final = response.get("access_token") @@ -1641,7 +1624,7 @@ async def get_generic_sso_response( discovery_document=discovery, response_convertor=response_convertor, ) - generic_sso: Final = SSOProvider( + generic_sso: Final[_GenericSSOClient] = SSOProvider( client_id=generic_client_id, client_secret=generic_client_secret, redirect_uri=redirect_url, @@ -1740,9 +1723,7 @@ async def get_generic_sso_response( access_token_str: str | None = combined_response.get("access_token") else: result = await _verify_generic_sso_or_fall_back_to_id_token( - generic_sso=cast( # cast-ok: fastapi-sso builds the provider class dynamically via create_provider - _UserinfoFallbackCapableSSO, generic_sso - ), + generic_sso=generic_sso, request=request, token_exchange_params=token_exchange_params, additional_headers=additional_generic_sso_headers_dict, @@ -1963,6 +1944,8 @@ async def get_user_info_from_db( ) return user_info + except _SSOUserAdmissionError: + raise except Exception as e: if resolved_by_database_only and PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): verbose_proxy_logger.exception("SSO login failed: the database could not resolve the account: %s", e) @@ -2328,7 +2311,6 @@ async def saml_callback(request: Request): general_settings, jwt_handler, master_key, - premium_user, prisma_client, saml_sso_cache, ) @@ -2352,8 +2334,6 @@ async def saml_callback(request: Request): user, retained = await SAMLAuthHandler.handle_acs(request=request, cache=saml_sso_cache, post_data=post_data) - await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) - ui_access_mode: Final = general_settings.get("ui_access_mode", None) relay_state: Final = post_data.get("RelayState") cp_return_to: Final[str | None] = ( @@ -2379,9 +2359,7 @@ async def _build_cli_sso_user_defined_values( result: OpenID | dict, parsed_openid_result: ParsedOpenIDResult, ) -> SSOUserDefinedValues | None: - from litellm.proxy.proxy_server import user_custom_sso - - custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso + custom_sso_handler: Final = _configured_custom_sso() user_id: Final = parsed_openid_result.get("user_id") if custom_sso_handler is not None: if inspect.iscoroutinefunction(custom_sso_handler): @@ -2467,10 +2445,8 @@ async def _complete_cli_sso_callback_session( ): from fastapi.responses import HTMLResponse - from litellm.proxy.proxy_server import user_custom_sso - user_id: Final = parsed_openid_result.get("user_id") - if user_custom_sso is None: + if _configured_custom_sso() is None: _require_sso_user_id(user_id) user_email: Final = parsed_openid_result.get("user_email") user_info: Final = await get_user_info_from_db( @@ -2821,12 +2797,20 @@ async def insert_sso_user( if result_openid and hasattr(result_openid, "provider"): new_user_request.metadata = {"auth_provider": getattr(result_openid, "provider")} - response: Final = await new_user( - data=new_user_request, - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) + from litellm.proxy.proxy_server import premium_user - return response + with free_sso_user_creation(new_user_request.user_id, enabled=premium_user is not True): + try: + return await new_user( + data=new_user_request, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + except ProxyException as error: + if premium_user is True: + raise + raise _SSOUserAdmissionError( + message=error.message, type=error.type, param=error.param, code=error.code + ) from error DefaultTeamJoin: TypeAlias = Callable[ @@ -3567,6 +3551,8 @@ class SSOAuthenticationHandler: prisma_client=prisma_client, ) return user_info + except _SSOUserAdmissionError: + raise except Exception as e: if ( user_info is None @@ -3826,7 +3812,6 @@ class SSOAuthenticationHandler: proxy_logging_obj, redis_usage_cache, user_api_key_cache, - user_custom_sso, ) from litellm.proxy.utils import get_prisma_client_or_throw from litellm.types.proxy.ui_sso import ReturnedUITokenObject @@ -3851,7 +3836,7 @@ class SSOAuthenticationHandler: # But if it is, we want their models preferences user_defined_values: SSOUserDefinedValues | None = None - custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso + custom_sso_handler: Final = _configured_custom_sso() if custom_sso_handler is not None: if inspect.iscoroutinefunction(custom_sso_handler): user_defined_values = await custom_sso_handler(result) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/__init__.py b/litellm/proxy/management_endpoints/usage_endpoints/__init__.py deleted file mode 100644 index 6e68dcd2a2e..00000000000 --- a/litellm/proxy/management_endpoints/usage_endpoints/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -""" -Usage endpoints package. - -Re-exports the router from endpoints module. -""" - -from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( # noqa: F401 - router, -) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py deleted file mode 100644 index 7a45d0438de..00000000000 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ /dev/null @@ -1,627 +0,0 @@ -""" -AI Usage Chat - uses LLM tool calling to answer questions about -usage/spend data by querying the aggregated daily activity endpoints. -""" - -import json -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence -from datetime import date -from typing import Final, Literal, NamedTuple, Protocol, cast, overload - -from fastapi import HTTPException -from typing_extensions import ReadOnly, TypedDict - -import litellm -from litellm._logging import verbose_proxy_logger -from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL -from litellm.proxy._types import CommonProxyErrors -from litellm.types.proxy.management_endpoints.common_daily_activity import ( - SpendAnalyticsPaginatedResponse, -) -from litellm.types.utils import ChatCompletionMessageToolCall - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- - -USAGE_AI_TEMPERATURE: Final = 0.2 - -TABLE_DAILY_USER_SPEND: Final = "litellm_dailyuserspend" -TABLE_DAILY_TEAM_SPEND: Final = "litellm_dailyteamspend" -TABLE_DAILY_TAG_SPEND: Final = "litellm_dailytagspend" - -ENTITY_FIELD_USER: Final = "user_id" -ENTITY_FIELD_TEAM: Final = "team_id" -ENTITY_FIELD_TAG: Final = "tag" - -PAGINATED_PAGE_SIZE: Final = 200 -MAX_CHAT_MESSAGES: Final = 20 -TOP_N_MODELS: Final = 15 -TOP_N_PROVIDERS: Final = 10 -TOP_N_KEYS: Final = 10 - -# --------------------------------------------------------------------------- -# Types -# --------------------------------------------------------------------------- - - -class SSEStatusEvent(TypedDict): - type: Literal["status"] - message: str - - -class SSEToolCallEvent(TypedDict, total=False): - type: Literal["tool_call"] - tool_name: str - tool_label: str - arguments: dict[str, str] - status: Literal["running", "complete", "error"] - error: str - - -class SSEChunkEvent(TypedDict): - type: Literal["chunk"] - content: str - - -class SSEDoneEvent(TypedDict): - type: Literal["done"] - - -class SSEErrorEvent(TypedDict): - type: Literal["error"] - message: str - - -SSEEvent = SSEStatusEvent | SSEToolCallEvent | SSEChunkEvent | SSEDoneEvent | SSEErrorEvent - - -class _EntityEntry(TypedDict, total=False): - metrics: ReadOnly[Mapping[str, float]] - metadata: ReadOnly[Mapping[str, str]] - - -class _DayDump(TypedDict, total=False): - breakdown: ReadOnly[Mapping[str, Mapping[str, _EntityEntry]]] - - -class _EntityTotal(NamedTuple): - """Running per-entity totals accumulated while summarising a usage dump.""" - - alias: str - spend: float - requests: float - tokens: float - - -class _UsageDump(Protocol): - @overload - def get(self, key: Literal["metadata"], default: Mapping[str, float], /) -> Mapping[str, float]: ... - @overload - def get(self, key: Literal["results"], default: Sequence[_DayDump], /) -> Sequence[_DayDump]: ... - - -class _ToolFunctionDef(TypedDict): - name: ReadOnly[str] - description: ReadOnly[str] - parameters: ReadOnly[Mapping[str, object]] - - -class _ToolDef(TypedDict): - type: ReadOnly[str] - function: ReadOnly[_ToolFunctionDef] - - -class ToolHandler(TypedDict): - fetch: Callable[..., Awaitable[_UsageDump]] - summarise: Callable[[_UsageDump], str] - label: str - - -# --------------------------------------------------------------------------- -# Tool definitions (OpenAI function-calling schema) -# --------------------------------------------------------------------------- - -_DATE_PARAMS: Final = { - "start_date": {"type": "string", "description": "Start date in YYYY-MM-DD format"}, - "end_date": {"type": "string", "description": "End date in YYYY-MM-DD format"}, -} - -_TOOL_USAGE: Final[_ToolDef] = { - "type": "function", - "function": { - "name": "get_usage_data", - "description": ( - "Fetch aggregated global usage/spend data. Returns daily spend, " - "token counts, request counts, and breakdowns by model, provider, " - "and API key. Use for overall spend, top models, top providers." - ), - "parameters": { - "type": "object", - "properties": { - **_DATE_PARAMS, - "user_id": { - "type": "string", - "description": "Optional user ID filter. Omit for global view.", - }, - }, - "required": ["start_date", "end_date"], - }, - }, -} - -_TOOL_TEAM: Final[_ToolDef] = { - "type": "function", - "function": { - "name": "get_team_usage_data", - "description": ( - "Fetch usage/spend data broken down by team. Use for questions " - "like 'which team spends the most' or 'show me team X usage'." - ), - "parameters": { - "type": "object", - "properties": { - **_DATE_PARAMS, - "team_ids": { - "type": "string", - "description": "Optional comma-separated team IDs. Omit for all teams.", - }, - }, - "required": ["start_date", "end_date"], - }, - }, -} - -_TOOL_TAG: Final[_ToolDef] = { - "type": "function", - "function": { - "name": "get_tag_usage_data", - "description": ( - "Fetch usage/spend data broken down by tag. Tags are labels " - "attached to requests (features, environments, credentials)." - ), - "parameters": { - "type": "object", - "properties": { - **_DATE_PARAMS, - "tags": { - "type": "string", - "description": "Optional comma-separated tag names. Omit for all tags.", - }, - }, - "required": ["start_date", "end_date"], - }, - }, -} - -TOOLS_BASE: Final = [_TOOL_USAGE] -TOOLS_ADMIN: Final = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG] - - -def get_tools_for_role(is_admin: bool) -> list[_ToolDef]: - """Return the tool list appropriate for the user's role.""" - return TOOLS_ADMIN if is_admin else TOOLS_BASE - - -_SYSTEM_PROMPT_BASE: Final = ( - "You are an AI assistant embedded in the LiteLLM Usage dashboard. " - "You help users understand their LLM API spend and usage data.\n\n" - "ALWAYS call the appropriate tool(s) first to fetch data before answering. " - "You may call multiple tools if the question spans different dimensions.\n\n" - "Guidelines:\n" - "- Be concise and specific. Use exact numbers from the data.\n" - "- Format costs as dollar amounts (e.g. $12.34).\n" - "- When comparing entities, show a ranked list.\n" - "- If data is empty or no results found, say so clearly.\n" - "- Do not hallucinate data — only use what the tools return.\n" - "- Today's date will be provided below. Use it to interpret relative dates " - "like 'this week', 'this month', 'last 7 days', etc." -) - -_TOOL_DESCRIPTIONS_ADMIN: Final = ( - "You have access to these tools:\n" - "- `get_usage_data`: Global/user-level usage (spend, models, providers, API keys)\n" - "- `get_team_usage_data`: Team-level usage breakdown\n" - "- `get_tag_usage_data`: Tag-level usage breakdown\n\n" -) - -_TOOL_DESCRIPTIONS_BASE: Final = ( - "You have access to this tool:\n- `get_usage_data`: Your usage data (spend, models, providers, API keys)\n\n" -) - - -def _build_system_prompt(is_admin: bool) -> str: - """Build role-appropriate system prompt with today's date.""" - tool_desc: Final = _TOOL_DESCRIPTIONS_ADMIN if is_admin else _TOOL_DESCRIPTIONS_BASE - return f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}Today's date: {date.today().isoformat()}" - - -# keep a public reference for test assertions -SYSTEM_PROMPT: Final = _SYSTEM_PROMPT_BASE - -# --------------------------------------------------------------------------- -# Data fetchers -# --------------------------------------------------------------------------- - - -def _parse_csv_ids(raw: str | None) -> list[str] | None: - if not raw: - return None - return [t.strip() for t in raw.split(",") if t.strip()] - - -async def _query_activity( - table_name: str, - entity_id_field: str, - entity_id: str | list[str] | None, - start_date: str, - end_date: str, - *, - use_aggregated: bool = False, -) -> SpendAnalyticsPaginatedResponse: - """Shared helper that calls the daily activity query layer.""" - from litellm.proxy.management_endpoints.common_daily_activity import ( - daily_activity_repository, - daily_activity_scope, - get_daily_activity, - get_daily_activity_aggregated, - ) - from litellm.proxy.proxy_server import prisma_client - - if use_aggregated: - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - repository: Final = daily_activity_repository(prisma_client) - scope: Final = daily_activity_scope( - table_name, - entity_id_field, - entity_id, - None, - None, - start_date, - end_date, - None, - None, - ) - return await get_daily_activity_aggregated( - repository, - scope, - ) - return await get_daily_activity( - prisma_client=prisma_client, - table_name=table_name, - entity_id_field=entity_id_field, - entity_id=entity_id, - entity_metadata_field=None, - start_date=start_date, - end_date=end_date, - model=None, - api_key=None, - page=1, - page_size=PAGINATED_PAGE_SIZE, - ) - - -async def _fetch_usage_data(start_date: str, end_date: str, user_id: str | None = None) -> _UsageDump: - resp: Final = await _query_activity( - TABLE_DAILY_USER_SPEND, - ENTITY_FIELD_USER, - user_id, - start_date, - end_date, - use_aggregated=True, - ) - return resp.model_dump(mode="json") - - -async def _fetch_team_usage_data(start_date: str, end_date: str, team_ids: str | None = None) -> _UsageDump: - resp: Final = await _query_activity( - TABLE_DAILY_TEAM_SPEND, - ENTITY_FIELD_TEAM, - _parse_csv_ids(team_ids), - start_date, - end_date, - ) - return resp.model_dump(mode="json") - - -async def _fetch_tag_usage_data(start_date: str, end_date: str, tags: str | None = None) -> _UsageDump: - resp: Final = await _query_activity( - TABLE_DAILY_TAG_SPEND, - ENTITY_FIELD_TAG, - _parse_csv_ids(tags), - start_date, - end_date, - ) - return resp.model_dump(mode="json") - - -# --------------------------------------------------------------------------- -# Summarisers — convert raw JSON to concise text the LLM can reason over -# --------------------------------------------------------------------------- - - -def _accumulate_breakdown( - results: Sequence[_DayDump], dimension: str, fields: Sequence[str] -) -> dict[str, dict[str, float]]: - """Aggregate a single breakdown dimension across days.""" - totals: Final[dict[str, dict[str, float]]] = {} - for day in results: - for key, entry in day.get("breakdown", {}).get(dimension, {}).items(): - if key not in totals: - totals[key] = {f: 0.0 for f in fields} - m = entry.get("metrics", {}) - for f in fields: - totals[key][f] += m.get(f, 0) - return totals - - -def _ranked_lines( - totals: dict[str, dict[str, float]], - fmt: Callable[[str, dict[str, float]], str], - limit: int, -) -> list[str]: - """Sort by spend descending, format each entry, and truncate.""" - return [fmt(name, vals) for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[:limit]] - - -def _summarise_usage_data(data: _UsageDump) -> str: - meta: Final = data.get("metadata", {}) - results: Final = data.get("results", []) - - header: Final = ( - f"Total Spend: ${meta.get('total_spend', 0):.4f}\n" - f"Total Requests: {meta.get('total_api_requests', 0)}\n" - f"Successful: {meta.get('total_successful_requests', 0)} | " - f"Failed: {meta.get('total_failed_requests', 0)}\n" - f"Total Tokens: {meta.get('total_tokens', 0)}" - ) - - models: Final = _accumulate_breakdown(results, "models", ["spend", "api_requests", "total_tokens"]) - providers: Final = _accumulate_breakdown(results, "providers", ["spend", "api_requests"]) - - model_lines: Final = _ranked_lines( - models, - lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)", - TOP_N_MODELS, - ) - provider_lines: Final = _ranked_lines( - providers, - lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs)", - TOP_N_PROVIDERS, - ) - - sections = [header, ""] - sections += ["Top Models by Spend:"] + (model_lines or [" (no data)"]) + [""] - sections += ["Top Providers by Spend:"] + (provider_lines or [" (no data)"]) - return "\n".join(sections) - - -def _summarise_entity_data(data: _UsageDump, entity_label: str) -> str: - """Summarise team/tag entity usage data.""" - results: Final = data.get("results", []) - if not results: - return f"No {entity_label} usage data found for the given date range." - - totals: Final[dict[str, _EntityTotal]] = {} - for day in results: - for eid, entry in day.get("breakdown", {}).get("entities", {}).items(): - previous = totals.get(eid) - m = entry.get("metrics", {}) - totals[eid] = _EntityTotal( - alias=previous.alias if previous is not None else entry.get("metadata", {}).get("alias", eid), - spend=(previous.spend if previous is not None else 0.0) + m.get("spend", 0), - requests=(previous.requests if previous is not None else 0) + m.get("api_requests", 0), - tokens=(previous.tokens if previous is not None else 0) + m.get("total_tokens", 0), - ) - - lines: Final = [f"{entity_label} Usage ({len(totals)} {entity_label.lower()}s):", ""] - for eid, d in sorted(totals.items(), key=lambda x: -x[1].spend): - label = d.alias if d.alias != eid else eid - lines.append(f"- {label} (ID: {eid}): ${d.spend:.4f} | {int(d.requests)} reqs | {int(d.tokens)} tokens") - return "\n".join(lines) - - -# --------------------------------------------------------------------------- -# Tool dispatch registry -# --------------------------------------------------------------------------- - -TOOL_HANDLERS: Final[dict[str, ToolHandler]] = { - "get_usage_data": ToolHandler( - fetch=_fetch_usage_data, - summarise=_summarise_usage_data, - label="global usage data", - ), - "get_team_usage_data": ToolHandler( - fetch=_fetch_team_usage_data, - summarise=lambda data: _summarise_entity_data(data, "Team"), - label="team usage data", - ), - "get_tag_usage_data": ToolHandler( - fetch=_fetch_tag_usage_data, - summarise=lambda data: _summarise_entity_data(data, "Tag"), - label="tag usage data", - ), -} - - -# --------------------------------------------------------------------------- -# SSE streaming -# --------------------------------------------------------------------------- - - -def _sse(event: SSEEvent) -> str: - return f"data: {json.dumps(event)}\n\n" - - -def _resolve_fetch_kwargs( - fn_name: str, - fn_args: Mapping[str, str], - user_id: str | None, - is_admin: bool, -) -> dict[str, str]: - """Build keyword arguments for a tool's fetch function.""" - start_date: Final = fn_args.get("start_date", "") - end_date: Final = fn_args.get("end_date", "") - if not start_date or not end_date: - raise ValueError("Missing required start_date or end_date from tool arguments") - kwargs: Final[dict[str, str]] = {"start_date": start_date, "end_date": end_date} - if fn_name == "get_usage_data": - if not is_admin: - if user_id is None: - # Defense-in-depth: the endpoint guard in usage_endpoints/endpoints.py - # should have already rejected this. If we ever reach here it means - # a future caller invoked the helper without scoping — fail loudly - # rather than issuing an unfiltered global query. - raise ValueError( - "Non-admin caller has user_id=None; refusing to issue an " - "unscoped query. Endpoint-level guard missing." - ) - kwargs["user_id"] = user_id - elif fn_args.get("user_id"): - kwargs["user_id"] = fn_args["user_id"] - elif fn_name == "get_team_usage_data" and fn_args.get("team_ids"): - kwargs["team_ids"] = fn_args["team_ids"] - elif fn_name == "get_tag_usage_data" and fn_args.get("tags"): - kwargs["tags"] = fn_args["tags"] - return kwargs - - -async def _execute_tool_call( - handler: ToolHandler, - fn_name: str, - fn_args: Mapping[str, str], - user_id: str | None, - is_admin: bool, -) -> str: - """Run a single tool and return the summarised result text, or the reason the query - layer rejected its dates so the model can correct them.""" - kwargs: Final = _resolve_fetch_kwargs(fn_name, fn_args, user_id, is_admin) - try: - raw_data: Final = await handler["fetch"](**kwargs) - except HTTPException as e: - reason: Final = _invalid_date_reason(e) - if reason is None: - raise - return f"Invalid date range: {reason}" - return handler["summarise"](raw_data) - - -def _invalid_date_reason(error: HTTPException) -> str | None: - """The reason in the daily activity query layer's 400 for invalid or reversed dates.""" - detail: Final = cast(object, error.detail) # cast-ok: starlette types detail as str; the query layer passes a dict - match detail: - case {"error": str() as reason} if error.status_code == 400: - return reason - case _: - return None - - -async def _process_tool_call( - tc: ChatCompletionMessageToolCall, - chat_messages: list[Mapping[str, object]], - user_id: str | None, - is_admin: bool, -) -> AsyncIterator[str]: - """Execute a single tool call, yielding SSE events for status.""" - fn_name: Final = tc.function.name - fn_args: Final[Mapping[str, str]] = json.loads(tc.function.arguments) - - allowed_names: Final = {t["function"]["name"] for t in get_tools_for_role(is_admin)} - handler: Final = TOOL_HANDLERS.get(fn_name) if fn_name is not None else None - - if fn_name is None or fn_name not in allowed_names or not handler: - chat_messages.append( - { - "role": "tool", - "tool_call_id": tc.id, - "content": f"Tool not available: {fn_name}", - } - ) - return - - tool_event_base: Final = { - "type": "tool_call", - "tool_name": fn_name, - "tool_label": handler["label"], - "arguments": fn_args, - } - yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "running"})) - - try: - tool_result = await _execute_tool_call(handler, fn_name, fn_args, user_id, is_admin) - yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "complete"})) - except Exception as e: - verbose_proxy_logger.error("Tool %s failed: %s", fn_name, e) - tool_result = f"Error fetching {handler['label']}. Please try again." - yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "error"})) - - chat_messages.append({"role": "tool", "tool_call_id": tc.id, "content": tool_result}) - - -async def _stream_final_response(model: str, chat_messages: list[Mapping[str, object]]) -> AsyncIterator[str]: - """Stream the final LLM response after tool results are appended.""" - yield _sse({"type": "status", "message": "Analyzing results..."}) - - response: Final = await litellm.acompletion( - model=model, - messages=chat_messages, - stream=True, - temperature=USAGE_AI_TEMPERATURE, - ) - async for chunk in response: - delta = chunk.choices[0].delta.content - if delta: - yield _sse({"type": "chunk", "content": delta}) - - -async def stream_usage_ai_chat( - messages: list[dict[str, str]], - model: str | None = None, - user_id: str | None = None, - is_admin: bool = False, -) -> AsyncGenerator[str, None]: - """Stream SSE events: status → tool_call → chunk → done.""" - resolved_model: Final = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL - truncated: Final = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages - chat_messages: Final[list[Mapping[str, object]]] = [ - {"role": "system", "content": _build_system_prompt(is_admin)}, - *truncated, - ] - - try: - yield _sse({"type": "status", "message": "Thinking..."}) - tools: Final = get_tools_for_role(is_admin) - response: Final = await litellm.acompletion( - model=resolved_model, - messages=chat_messages, - tools=tools, - temperature=USAGE_AI_TEMPERATURE, - ) - choice: Final = response.choices[0] - - if not choice.message.tool_calls: - if choice.message.content: - yield _sse({"type": "chunk", "content": choice.message.content}) - yield _sse({"type": "done"}) - return - - chat_messages.append(choice.message.model_dump()) - for tc in choice.message.tool_calls: - async for event in _process_tool_call(tc, chat_messages, user_id, is_admin): - yield event - async for event in _stream_final_response(resolved_model, chat_messages): - yield event - yield _sse({"type": "done"}) - - except Exception as e: - verbose_proxy_logger.error("AI usage chat failed: %s", e) - yield _sse( - { - "type": "error", - "message": "An internal error occurred. Please try again.", - } - ) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py deleted file mode 100644 index cff5bef05a0..00000000000 --- a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py +++ /dev/null @@ -1,77 +0,0 @@ -""" -USAGE AI CHAT ENDPOINTS - -/usage/ai/chat - Stream AI chat responses about usage data -""" - -from typing import Final, Literal - -from fastapi import APIRouter, Depends, Request -from fastapi.responses import StreamingResponse -from pydantic import Field - -import litellm -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.sse_keepalive import ( - SSE_COMMENT_PING, - wrap_sse_stream_with_keepalive_pings, -) -from litellm.types.llms.base import LiteLLMBaseModel - -router: Final = APIRouter() - - -class ChatMessage(LiteLLMBaseModel): - role: Literal["user", "assistant"] - content: str - - -class UsageAIChatRequest(LiteLLMBaseModel): - messages: list[ChatMessage] = Field(..., description="Chat messages (user/assistant history)") - model: str | None = Field(default=None, description="Model to use for AI chat") - - -@router.post( - "/usage/ai/chat", - tags=["Budget & Spend Tracking"], - dependencies=[Depends(user_api_key_auth)], -) -async def usage_ai_chat( - data: UsageAIChatRequest, - request: Request, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - AI chat about usage data. Streams SSE events with the AI response. - The AI agent has access to tools that query aggregated daily activity data. - """ - from litellm.proxy.management_endpoints.common_utils import ( - require_caller_user_id_for_non_admin, - user_api_key_has_admin_view, - ) - from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( - stream_usage_ai_chat, - ) - - is_admin: Final = user_api_key_has_admin_view(user_api_key_dict) - if is_admin: - user_id = user_api_key_dict.user_id - else: - user_id = require_caller_user_id_for_non_admin(user_api_key_dict) - messages: Final = [{"role": m.role, "content": m.content} for m in data.messages] - - return StreamingResponse( - wrap_sse_stream_with_keepalive_pings( - stream_usage_ai_chat( - messages=messages, - model=data.model, - user_id=user_id, - is_admin=is_admin, - ), - ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, - ping_chunk=SSE_COMMENT_PING, - ), - media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, - ) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 3468396c6f6..920970427ba 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -310,7 +310,7 @@ async def _resolve_mcp_server_identifiers_to_ids( identifiers=identifiers, prisma_client=prisma_client, ): - server_id = getattr(server, "server_id", None) + server_id = server.server_id if not server_id: continue for identifier in identifiers: diff --git a/litellm/proxy/management_helpers/team_list_sorting.py b/litellm/proxy/management_helpers/team_list_sorting.py new file mode 100644 index 00000000000..caddd8afc08 --- /dev/null +++ b/litellm/proxy/management_helpers/team_list_sorting.py @@ -0,0 +1,109 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +from litellm.proxy.utils import PrismaClient + +COMPUTED_TEAM_SORTS: Final = MappingProxyType( + { + "organization_alias": "COALESCE((SELECT NULLIF(o.organization_alias, '') " + 'FROM "LiteLLM_OrganizationTable" o WHERE o.organization_id = t.organization_id), ' + "NULLIF(t.organization_id, ''))", + "members": "CASE WHEN jsonb_typeof(t.members_with_roles) = 'array' " + "THEN jsonb_array_length(t.members_with_roles) ELSE 0 END", + "models": "cardinality(t.models)", + "keys_count": '(SELECT COUNT(*) FROM "LiteLLM_VerificationToken" k WHERE k.team_id = t.team_id)', + } +) +_WHERE: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class UnsupportedTeamFilter: + field: str + + +def _field_sql(field: str, value: JsonValue, index: int) -> tuple[str, tuple[object, ...]] | UnsupportedTeamFilter: + if field not in ("team_id", "team_alias", "organization_id"): + return UnsupportedTeamFilter(field) + column: Final = f't."{field}"' + match value: + case str(text): + return f"{column} = ${index}", (text,) + case {"in": list(values)}: + return f"{column} = ANY(${index}::text[])", (values,) + case {"contains": str(text), "mode": "insensitive"}: + return f"{column} ILIKE ${index}", (f"%{text}%",) + case {"startsWith": str(text)}: + return f"{column} LIKE ${index}", (f"{text}%",) + case _: + return UnsupportedTeamFilter(field) + + +def _group_tokens( + field: str, value: Sequence[JsonValue] | Mapping[str, JsonValue] +) -> tuple[str | tuple[str, JsonValue], ...]: + children: Final = tuple(value.items()) if isinstance(value, Mapping) else tuple(("AND", item) for item in value) + if not children: + return ("FALSE" if field == "OR" else "TRUE",) + return ( + "(", + *chain.from_iterable(((f" {field} " if index else ""), child) for index, child in enumerate(children)), + ")", + ) + + +def _where_sql(where: Mapping[str, JsonValue]) -> tuple[str, tuple[object, ...]] | UnsupportedTeamFilter: + pending: Final[list[str | tuple[str, JsonValue]]] = [("AND", dict(where))] # mutable-ok: iterative tree traversal + fragments: Final[list[str]] = [] # mutable-ok: append-only SQL output avoids quadratic tuple copying + params: Final[list[object]] = [] # mutable-ok: append-only bindings follow traversal order + while pending: + match pending.pop(): + case str(fragment): + fragments.append(fragment) + case (field, value) if field in ("AND", "OR") and isinstance(value, (list, dict)): + pending.extend(reversed(_group_tokens(field, value))) + case (field, value): + match _field_sql(field, value, len(params) + 1): + case UnsupportedTeamFilter() as failure: + return failure + case (fragment, values): + fragments.append(fragment) + params.extend(values) + return "".join(fragments), tuple(params) + + +class _TeamId(BaseModel): + model_config = ConfigDict(frozen=True) + team_id: str + + +_TEAM_IDS: Final = TypeAdapter(tuple[_TeamId, ...]) + + +async def sorted_team_ids( + prisma_client: PrismaClient, + *, + where: Mapping[str, object], + sort_by: str, + descending: bool, + skip: int, + take: int, +) -> Sequence[str] | UnsupportedTeamFilter: + result: Final = _where_sql(_WHERE.validate_python(where)) + if isinstance(result, UnsupportedTeamFilter): + return result + predicate, params = result + expression: Final = COMPUTED_TEAM_SORTS[sort_by] + direction: Final = "DESC" if descending else "ASC" + sql: Final = ( + f'SELECT t.team_id FROM "LiteLLM_TeamTable" t WHERE {predicate} ' + f"ORDER BY {expression} {direction} NULLS LAST, t.team_id ASC " + f"LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" + ) + rows: Final = _TEAM_IDS.validate_python(await prisma_client.db.query_raw(sql, *params, take, skip)) + return tuple(row.team_id for row in rows) diff --git a/litellm/proxy/management_helpers/user_list_sorting.py b/litellm/proxy/management_helpers/user_list_sorting.py new file mode 100644 index 00000000000..d3d2e323cb0 --- /dev/null +++ b/litellm/proxy/management_helpers/user_list_sorting.py @@ -0,0 +1,77 @@ +from collections.abc import Sequence +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy.utils import PrismaClient + + +class _UserId(BaseModel): + model_config = ConfigDict(frozen=True) + user_id: str + + +_USER_IDS: Final = TypeAdapter(tuple[_UserId, ...]) + + +async def sorted_user_ids( + prisma_client: PrismaClient, + *, + sort_by: Literal["status", "key_count"], + descending: bool, + skip: int, + take: int, + role: str | None, + user_ids: str | None, + user_email: str | None, + team: str | None, + sso_user_ids: str | None, + organization_ids: str | None, + search: str | None, +) -> Sequence[str]: + ids: Final = tuple(value.strip() for value in (user_ids or "").split(",") if value.strip()) + sso_ids: Final = tuple(value.strip() for value in (sso_user_ids or "").split(",") if value.strip()) + org_ids: Final = tuple(value.strip() for value in (organization_ids or "").split(",") if value.strip()) + expression: Final = ( + "CASE WHEN u.metadata->'scim_active' = 'false'::jsonb THEN 1 ELSE 0 END" + if sort_by == "status" + else '(SELECT COUNT(*) FROM "LiteLLM_VerificationToken" k ' + "WHERE k.user_id = u.user_id AND k.team_id IS DISTINCT FROM $11::text)" + ) + direction: Final = "DESC" if descending else "ASC" + sql: Final = f""" + SELECT u.user_id FROM "LiteLLM_UserTable" u + WHERE ($1::text IS NULL OR u.user_role = $1) + AND ($2::text[] IS NULL + OR CASE WHEN cardinality($2::text[]) = 1 + THEN u.user_id ILIKE '%' || ($2::text[])[1] || '%' + ELSE u.user_id = ANY($2::text[]) END) + AND ($3::text IS NULL OR u.user_email ILIKE '%' || $3 || '%') + AND ($4::text IS NULL OR $4 = ANY(u.teams)) + AND ($5::boolean OR u.sso_user_id = ANY($6::text[])) + AND (cardinality($7::text[]) = 0 OR EXISTS ( + SELECT 1 FROM "LiteLLM_OrganizationMembership" m + WHERE m.user_id = u.user_id AND m.organization_id = ANY($7::text[]))) + AND ($8::text IS NULL OR u.user_id ILIKE '%' || $8 || '%' + OR u.user_email ILIKE '%' || $8 || '%') + ORDER BY {expression} {direction}, u.user_id ASC + LIMIT $9 OFFSET $10 + """ + rows: Final = _USER_IDS.validate_python( + await prisma_client.db.query_raw( + sql, + role or None, + ids if user_ids else None, + user_email, + team, + sso_user_ids is None, + sso_ids, + org_ids, + search or None, + take, + skip, + *((UI_SESSION_TOKEN_TEAM_ID,) if sort_by == "key_count" else ()), + ) + ) + return tuple(row.user_id for row in rows) diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 34c80bcdce3..4457df460d4 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -181,8 +181,8 @@ async def _assert_write_access( """ if _is_admin(user_api_key_dict): return - row_user_id: Final = getattr(row, "user_id", None) - row_team_id: Final = getattr(row, "team_id", None) + row_user_id: Final = row.user_id + row_team_id: Final = row.team_id # Personal ownership. if row_user_id and row_user_id == user_api_key_dict.user_id: diff --git a/litellm/proxy/moyai_endpoints.py b/litellm/proxy/moyai_endpoints.py index c4d8ad139f7..45f4502ec4e 100644 --- a/litellm/proxy/moyai_endpoints.py +++ b/litellm/proxy/moyai_endpoints.py @@ -20,9 +20,10 @@ from typing import TYPE_CHECKING, Annotated, Final from urllib.parse import urlencode, urlparse from fastapi import APIRouter, Depends, HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm._internal_context import with_service_target +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( @@ -114,6 +115,9 @@ def _sign_connect_code(master_key: str, moyai_url: str, user_id: str | None) -> return f"{_b64url(payload)}.{_b64url(signature)}" +_DECODED_CONNECT_PAYLOAD: Final = TypeAdapter(object) + + def _decode_connect_code(master_key: str, code: str) -> _MoyaiConnectCode: try: payload_b64, signature_b64 = code.split(".", 1) @@ -125,18 +129,23 @@ def _decode_connect_code(master_key: str, code: str) -> _MoyaiConnectCode: if not hmac.compare_digest(signature, expected): raise HTTPException(status_code=400, detail="Invalid Moyai connect code") try: - payload: Final = json.loads(payload_raw) + payload: Final = as_str_mapping(_DECODED_CONNECT_PAYLOAD.validate_python(json.loads(payload_raw))) except (ValueError, TypeError): raise HTTPException(status_code=400, detail="Invalid Moyai connect code") - if not isinstance(payload, dict) or not isinstance(payload.get("exp"), int) or payload["exp"] < int(time.time()): + if payload is None: raise HTTPException(status_code=400, detail="Invalid Moyai connect code") - if not isinstance(payload.get("moyai_origin"), str) or not isinstance(payload.get("nonce"), str): + exp: Final = payload.get("exp") + if not isinstance(exp, int) or exp < int(time.time()): + raise HTTPException(status_code=400, detail="Invalid Moyai connect code") + moyai_origin: Final = payload.get("moyai_origin") + nonce: Final = payload.get("nonce") + if not isinstance(moyai_origin, str) or not isinstance(nonce, str): raise HTTPException(status_code=400, detail="Invalid Moyai connect code") user_id: Final = payload.get("user_id") return _MoyaiConnectCode( - moyai_origin=payload["moyai_origin"], - nonce=payload["nonce"], - exp=payload["exp"], + moyai_origin=moyai_origin, + nonce=nonce, + exp=exp, user_id=user_id if isinstance(user_id, str) else None, ) diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index 1aa24eed4cd..5d596aece32 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -4,12 +4,15 @@ from dataclasses import dataclass from itertools import chain from typing import BinaryIO, Final, NoReturn +from pydantic import TypeAdapter from typing_extensions import assert_never +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import ProxyException from litellm.proxy.openai_files_endpoints.file_usage_caps import FileUsageLimit, describe_limit_source _MB: Final = 1024 * 1024 +_DECODED_LINE: Final = TypeAdapter(object) @dataclass(frozen=True, slots=True) @@ -95,10 +98,10 @@ def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]: def _check_line(line_number: int, raw_line: bytes, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: try: - parsed: Final = json.loads(raw_line) + parsed: Final = as_str_mapping(_DECODED_LINE.validate_python(json.loads(raw_line))) except (ValueError, RecursionError): return BatchFileInvalidJsonLine(line_number=line_number) - if not isinstance(parsed, dict): + if parsed is None: return BatchFileLineNotObject(line_number=line_number) missing: Final = next((key for key in line_shape.required_keys if key not in parsed), None) if missing is None: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 2e07694876b..f7289445d3a 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -477,6 +477,18 @@ async def get_authorized_credentials_for_model( ) +def deployment_drop_params( + llm_router: Optional["Router"], model_id: str, request: Mapping[str, object] +) -> Mapping[str, bool | str]: + deployment: Final = ( + llm_router.get_credential_deployment(model_id=model_id) + if llm_router is not None and "drop_params" not in request + else None + ) + drop_params: Final = deployment.litellm_params.drop_params if deployment is not None else None + return MappingProxyType({}) if drop_params is None else MappingProxyType({"drop_params": drop_params}) + + def _key_allowlist_covers(public_model_name: str, key_model_allowlist: Sequence[str]) -> bool: from litellm.proxy.auth.auth_checks import model_matches_any_wildcard_pattern_in_list diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 8823f01d781..d2ceefffe47 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -76,6 +76,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # add_internal_model_credentials, apply_team_provider_credentials, authorize_model_for_key, + deployment_drop_params, encode_file_id_with_model, extract_file_creation_params, get_authorized_credentials_for_model, @@ -457,7 +458,7 @@ async def route_create_file( # Merge credentials into the request prepare_data_with_credentials( data=_create_file_request, - credentials=credentials, + credentials={**credentials, **deployment_drop_params(llm_router, model, _create_file_request)}, ) # Create the file with model credentials diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index afc1e1c96fa..1f507982e65 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -8,10 +8,11 @@ storage backends (e.g., Azure Blob Storage) and managing associated metadata. import base64 import time from collections.abc import Sequence -from typing import Any, Final, cast +from typing import Final from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid as uuid_module +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend @@ -19,6 +20,7 @@ from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.types.llms.openai import OpenAIFileObject, OpenAIFilesPurpose +from litellm.types.passthrough_endpoints.managed_id_rewriter import ManagedFileIdWriter from litellm.types.utils import ExtractedFileData, SpecialEnums @@ -232,11 +234,10 @@ class StorageBackendFileService: proxy_logging_obj: Proxy logging object user_api_key_dict: User API key authentication data """ - managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") + managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if not managed_files_obj or not isinstance(managed_files_obj, BaseFileEndpoints): verbose_proxy_logger.warning("Managed files hook not available, skipping managed files storage") return - managed_files_obj = cast(Any, managed_files_obj) # Create model mappings using storage URL model_mappings: Final = {model_name: storage_url for model_name in target_model_names} @@ -260,7 +261,10 @@ class StorageBackendFileService: ) # Store in managed files - await managed_files_obj.store_unified_file_id( + managed_files_writer: Final = require_attributes( + managed_files_obj, ManagedFileIdWriter, "store_unified_file_id" + ) + await managed_files_writer.store_unified_file_id( file_id=base64_unified_file_id, file_object=file_object, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, diff --git a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py new file mode 100644 index 00000000000..3087faa9f67 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py @@ -0,0 +1,87 @@ +import hmac +from types import MappingProxyType +from typing import Annotated, Final, Literal + +from fastapi import APIRouter, Depends, HTTPException, Request, Response + +from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import pass_through_request + +router: Final = APIRouter() +_HEADERS: Final = frozenset( + { + "authorization", + "chatgpt-account-id", + "content-type", + "accept", + "user-agent", + "originator", + "openai-beta", + "session-id", + "thread-id", + "session_id", + "conversation_id", + "x-codex-turn-metadata", + "x-codex-turn-state", + } +) + + +@router.get("/backend-api/codex/{endpoint}", tags=["ChatGPT Pass-through", "pass-through"]) +@router.post( + "/backend-api/codex/{endpoint}", + tags=["ChatGPT Pass-through", "pass-through"], + openapi_extra={ + "requestBody": { + "required": True, + "content": {"application/json": {"schema": {"type": "object", "additionalProperties": True}}}, + } + }, + responses={ + 200: { + "description": "Unmodified Codex response or server-sent events", + "content": { + "application/json": {"schema": {"type": "object", "additionalProperties": True}}, + "text/event-stream": {"schema": {"type": "string"}}, + }, + } + }, +) +async def codex_passthrough_route( + endpoint: Literal["models", "responses"], + request: Request, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> Response: + """Forward a client's ChatGPT subscription with separate LiteLLM authentication.""" + expected_method: Final = "GET" if endpoint == "models" else "POST" + if request.method != expected_method: + raise HTTPException(status_code=405, detail="Method not allowed", headers={"Allow": expected_method}) + + authorization: Final = request.headers.get("authorization", "") + scheme, _, token = authorization.partition(" ") + gateway_key: Final = request.headers.get("x-litellm-api-key", "") + gateway_parts: Final = gateway_key.split() + gateway_token: Final = ( + gateway_parts[-1] + if len(gateway_parts) == 1 or (len(gateway_parts) == 2 and gateway_parts[0].lower() == "bearer") + else "" + ) + if scheme.lower() != "bearer" or not token.strip() or not gateway_token: + raise HTTPException( + status_code=401, + detail="Send the ChatGPT bearer in Authorization and the LiteLLM key in x-litellm-api-key", + ) + if hmac.compare_digest(token.strip().encode(), gateway_token.encode()): + raise HTTPException(status_code=401, detail="ChatGPT and LiteLLM credentials must be separate") + + headers: Final = MappingProxyType({name: value for name, value in request.headers.items() if name in _HEADERS}) + return await pass_through_request( + request=request, + target=f"{CHATGPT_API_BASE}/{endpoint}", + custom_headers=dict(headers), + user_api_key_dict=user_api_key_dict, + custom_llm_provider="chatgpt", + upstream_headers_filter=lambda _: headers, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 16fbd1de376..9eef4f0ef9d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1305,7 +1305,7 @@ async def handle_bedrock_count_tokens( verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e) from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers - provider_headers: Final = getattr(getattr(e, "response", None), "headers", None) + provider_headers: Final = e.response.headers raise HTTPException( status_code=e.status_code, detail={"error": e.message}, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a7f3a4f84a7..15f70073c7a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -1,23 +1,26 @@ import asyncio import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ANTHROPIC_BATCHES_ROUTE -from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import ( + native_messages_reported_cost, + override_stamped_response_cost, + response_cost_in_hidden_params, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, ) from litellm.llms.anthropic import get_anthropic_config -from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception from litellm.llms.anthropic.pass_through.stream_assembly import ( build_complete_streaming_response, build_complete_streaming_response_legacy, @@ -55,6 +58,38 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +def _event_payload(event: str) -> Mapping[str, object] | None: + data_lines: Final = tuple( + line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:") + ) + if not data_lines or data_lines[0] in ("", "[DONE]"): + return None + try: + return _JSON_OBJECT.validate_json(data_lines[0]) + except ValidationError: + return None + + +def _message_delta_usages(all_chunks: Sequence[str | bytes]) -> Iterator[object]: + for chunk in reversed(all_chunks): + yield from ( + payload.get("usage") + for payload in map(_event_payload, reversed(split_sse_chunk_into_events(chunk))) + if payload is not None and payload.get("type") == "message_delta" + ) + + +def _stamp_reported_stream_cost( + response: ModelResponse, all_chunks: Sequence[str | bytes], logging_obj: LiteLLMLoggingObj +) -> None: + """The final ``message_delta`` is where a provider that prices its own responses reports the cost.""" + provider: Final = logging_obj.model_call_details.get("custom_llm_provider") + costs: Final = (native_messages_reported_cost(usage, provider) for usage in _message_delta_usages(all_chunks)) + reported: Final = next((cost for cost in costs if cost is not None), None) + if reported is not None: + override_stamped_response_cost(response, reported) + + class AnthropicPassthroughLoggingHandler: @staticmethod def anthropic_passthrough_handler( @@ -224,6 +259,7 @@ class AnthropicPassthroughLoggingHandler: ) if not isinstance(response, ModelResponse): return response + _stamp_reported_stream_cost(response, all_chunks, litellm_logging_obj) recovered_usage: Final = AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( response=response, all_chunks=all_chunks, model=model ) @@ -256,26 +292,6 @@ class AnthropicPassthroughLoggingHandler: ), ) - @staticmethod - async def log_error_frame_as_failure( - attempt_logging_obj: LiteLLMLoggingObj, - request_body: Mapping[str, object], - all_chunks: Sequence[str | bytes], - error_event: tuple[str, str, int], - ) -> None: - """A stream closed at the provider's `event: error` frame is one failed attempt of the request: it logs - through the attempt-scoped copy its caller took at the close, so the usage it consumed never reaches - the success log of the attempt the router opens in its place.""" - error_type, message, status_code = error_event - await asyncify(AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure)( - litellm_logging_obj=attempt_logging_obj, request_body=request_body, all_chunks=all_chunks - ) - await attempt_logging_obj.dispatch_failure_handlers( - anthropic_error_frame_exception(error_type, message, status_code, str(attempt_logging_obj.model)), - "", - prefer_async_handlers=True, - ) - @staticmethod def _cost_partial_stream_or_zero( partial_response: ModelResponse | TextCompletionResponse, model: str, logging_obj: LiteLLMLoggingObj @@ -308,13 +324,21 @@ class AnthropicPassthroughLoggingHandler: if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/") else model ) + custom_pricing: Final = use_custom_pricing_for_model( + litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None) + ) + reported_cost: Final = ( + response_cost_in_hidden_params(litellm_model_response) + if isinstance(litellm_model_response, ModelResponse) + else None + ) + if reported_cost is not None and not custom_pricing: + return reported_cost return litellm.completion_cost( completion_response=litellm_model_response, model=model_for_cost, custom_llm_provider=custom_llm_provider, - custom_pricing=use_custom_pricing_for_model( - litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None) - ), + custom_pricing=custom_pricing, router_model_id=logging_obj.get_router_model_id(), ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py index 7527c504c59..3fb4fbd555d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py @@ -27,11 +27,15 @@ COMPREHEND_MEDICAL_COST_PER_UNIT_USD: Final[Mapping[str, float]] = MappingProxyT COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS: Final = frozenset(COMPREHEND_MEDICAL_COST_PER_UNIT_USD) +def _operation_from_headers(headers: Mapping[str, str]) -> str: + target: Final = headers.get("x-amz-target", "") + return target.split(".")[-1] + + class ComprehendMedicalPassthroughLoggingHandler: @staticmethod def _operation_from_response(httpx_response: httpx.Response) -> str: - target: Final[str] = httpx_response.request.headers.get("x-amz-target", "") - return target.split(".")[-1] + return _operation_from_headers(httpx_response.request.headers) @staticmethod def get_cost_for_operation(operation: str, text: str) -> float: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index c763b5dc26e..c410150df70 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -10,6 +10,7 @@ from typing import Final, Literal, TypeAlias from urllib.parse import urlparse import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -95,6 +96,8 @@ def is_openai_compatible_url(url_route: str | None) -> bool: hostname: Final = parsed_url.hostname if not hostname: return False + if hostname == "chatgpt.com" and parsed_url.path.startswith("/backend-api/codex/"): + return True if _hostname_matches(hostname, _OPENAI_HOSTNAMES): return True if _hostname_matches(hostname, _AZURE_OPENAI_HOSTNAMES): @@ -198,7 +201,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url: Final = urlparse(url_route) - return _is_openai_compatible_host(parsed_url.hostname) and ( + return is_openai_compatible_url(url_route) and ( "/v1/responses" in parsed_url.path or "/responses" in parsed_url.path ) @@ -424,7 +427,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) = None handler_instance: Final = OpenAIPassthroughLoggingHandler() - custom_llm_provider: Final = kwargs.get("custom_llm_provider", "openai") + custom_llm_provider: Final = TypeAdapter(str).validate_python( + kwargs.get("custom_llm_provider", logging_obj.model_call_details.get("custom_llm_provider", "openai")) + ) if is_chat_completions: # Handle chat completions with existing logic diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py index e277655f1e1..2afc1d17bc2 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py @@ -36,6 +36,7 @@ from litellm.types.passthrough_endpoints.tinyfish import ( from litellm.types.utils import StandardPassThroughResponseObject _RUN_ADAPTER: Final = TypeAdapter(TinyfishRun) +_DECODED_JSON: Final = TypeAdapter(object) _EMPTY_KWARGS: Final[Mapping[str, object]] = MappingProxyType({}) @@ -323,7 +324,7 @@ class TinyFishPassthroughLoggingHandler: "TinyFish passthrough: GET /v1/runs/%s returned %s", safe_run_id, response.status_code ) return None - payload: Final[object] = response.json() # any-ok: httpx Response.json() -> Any + payload: Final = _DECODED_JSON.validate_python(response.json()) return _parse_run(payload) except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request verbose_proxy_logger.warning("[Non blocking logging error] TinyFish run fetch failed: %s", e) @@ -413,7 +414,7 @@ def _run_id_from_sse_chunks(all_chunks: Sequence[str]) -> str | None: if not line.startswith("data:"): continue try: - event_payload: object = json.loads(line[5:].strip()) # any-ok: json.loads -> Any + event_payload: object = _DECODED_JSON.validate_python(json.loads(line[5:].strip())) except json.JSONDecodeError: continue event = _parse_run(event_payload) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py index 79f6dee90f7..5b45316a44c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py @@ -492,6 +492,10 @@ def s3_media_url(media_uri: str, aws_region_name: str) -> str | None: return f"https://{bucket}.s3.{aws_region_name}.{dns_suffix}/{quote(key)}" +def _header_or_default(headers: Mapping[str, str], name: str, default: str) -> str: + return headers.get(name, default) + + def media_predates_job(headers: Mapping[str, str], job_created_at: float) -> bool: try: modified_at: Final = parsedate_to_datetime(headers["last-modified"]).timestamp() @@ -501,7 +505,7 @@ def media_predates_job(headers: Mapping[str, str], job_created_at: float) -> boo async def write_media_within_limit(response: httpx.Response, media_file: IO[bytes], max_bytes: int) -> bool: - if int(response.headers.get("content-length", "0")) > max_bytes: + if int(_header_or_default(response.headers, "content-length", "0")) > max_bytes: return False async for chunk in response.aiter_bytes(): _ = media_file.write(chunk) @@ -591,8 +595,7 @@ class TranscribePassthroughLoggingHandler: @staticmethod def _operation_from_response(httpx_response: httpx.Response) -> str: - headers: Final[Mapping[str, str]] = httpx_response.request.headers - target: Final[str] = headers.get("x-amz-target", "") + target: Final = _header_or_default(httpx_response.request.headers, "x-amz-target", "") return target.split(".")[-1] @staticmethod diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fe7ac30e434..fc0e5be3f4c 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -118,6 +118,7 @@ from litellm.proxy.litellm_pre_call_utils import ( # noqa: F401 # legacy modul _key_or_team_allows_client_pricing_override, # pyright: ignore[reportPrivateUsage] # reuse the proxy's pricing trust policy _strip_client_pricing_overrides, # pyright: ignore[reportPrivateUsage] # sanitize before trusted hooks add guardrail costs get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above + redact_credential_headers, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.route_priority import ( @@ -409,6 +410,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): @staticmethod def get_endpoint_type(url: str, custom_llm_provider: str | None = None) -> EndpointType: parsed_url: Final = urlparse(url) + if custom_llm_provider == "chatgpt": + return EndpointType.OPENAI if custom_llm_provider in ("typesafe", "laya", "bespoke") and parsed_url.path.removesuffix("/").endswith( "/v1/systemone" ): @@ -723,7 +726,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "url": str(request.url), "method": request.method, "body": copy.copy(_parsed_body), # use copy instead of deepcopy - "headers": request.headers, + "headers": redact_credential_headers(request.headers), }, }, "call_type": "pass_through_endpoint", @@ -880,8 +883,8 @@ def _build_passthrough_failure_request_payload( @dataclass(frozen=True, slots=True) class _TeamCallbackWiring: - success_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg - failure_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg + success_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg + failure_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg logging_kwargs: dict[str, str | dict[str, str]] | None = None # mutable-ok: Logging.__init__ arg @@ -1490,7 +1493,7 @@ async def pass_through_request( additional_args={ "complete_input_dict": _parsed_body, "api_base": str(logging_url), - "headers": upstream_headers, + "headers": redact_credential_headers(upstream_headers), }, ) stream = HttpPassThroughEndpointHelpers.update_stream_param_based_on_request_body( @@ -1561,7 +1564,7 @@ async def pass_through_request( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, response=relay_response, - request_headers=dict(request.headers), + request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: _response_headers.update(callback_headers) @@ -1653,7 +1656,7 @@ async def pass_through_request( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, response=detected_relay_response, - request_headers=dict(request.headers), + request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: _response_headers.update(callback_headers) @@ -1703,7 +1706,7 @@ async def pass_through_request( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, response=response, - request_headers=dict(request.headers), + request_headers=dict(redact_credential_headers(request.headers)), ) if relay_callback_headers: relay_custom_headers.update(relay_callback_headers) @@ -1878,7 +1881,7 @@ async def pass_through_request( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, response=response, - request_headers=dict(request.headers), + request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: custom_headers.update(callback_headers) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 1fcfe2e279a..b6a067b7666 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,5 +1,5 @@ import traceback -from collections.abc import Coroutine, Sequence +from collections.abc import Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from typing import Final, Protocol @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks 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_frame_exception from litellm.proxy._types import PassThroughEndpointLoggingResultValues from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames @@ -87,13 +88,32 @@ class PassThroughStreamingHandler: ) try: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=litellm_logging_obj.dispatch_failure_handlers( - exception, traceback.format_exc(), prefer_async_handlers=True - ) + async_coroutine=litellm_logging_obj.dispatch_attempt_failure_handlers(exception, traceback.format_exc()) ) except Exception as e: verbose_proxy_logger.error("Error scheduling stream failure logging: %s", e) + @staticmethod + async def _schedule_stream_end_logging( + litellm_logging_obj: LiteLLMLoggingObj, + endpoint_type: EndpointType, + request_body: Mapping[str, object], + raw_bytes: Sequence[bytes], + error_event: tuple[str, str, int] | None, + success_logging: Callable[[], Coroutine[None, None, None]], + ) -> None: + if error_event is None: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=success_logging()) + return + error_type, message, status_code = error_event + await PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=litellm_logging_obj, + endpoint_type=endpoint_type, + request_body=dict(request_body), + raw_bytes=raw_bytes, + exception=anthropic_error_frame_exception(error_type, message, status_code, str(litellm_logging_obj.model)), + ) + @staticmethod def _record_partial_usage_for_failure( litellm_logging_obj: LiteLLMLoggingObj, @@ -276,17 +296,13 @@ class PassThroughStreamingHandler: else PassThroughStreamingHandler._error_frame_closed_at(endpoint_type, raw_bytes) ) try: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=( - _build_logging_coroutine() - if superseded_error_event is None - else AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( - attempt_logging_obj=litellm_logging_obj.attempt_scoped_copy(), - request_body=resolved_request_body, - all_chunks=tuple(raw_bytes), - error_event=superseded_error_event, - ) - ) + await PassThroughStreamingHandler._schedule_stream_end_logging( + litellm_logging_obj=litellm_logging_obj, + endpoint_type=endpoint_type, + request_body=resolved_request_body, + raw_bytes=tuple(raw_bytes), + error_event=superseded_error_event, + success_logging=_build_logging_coroutine, ) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) @@ -295,10 +311,6 @@ class PassThroughStreamingHandler: @staticmethod def _error_frame_closed_at(endpoint_type: EndpointType, raw_bytes: Sequence[bytes]) -> tuple[str, str, int] | None: - """The provider `event: error` frame of a stream the router closed because a retry or fallback took the - request over. That attempt is logged as its own failure, on a copy taken now so the retry's rebinding of - the shared logging object never reaches it. A stream that ran to its end or that the client closed keeps - logging through the success route, so its usage is still billed.""" from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( closing_superseded_attempt, parse_anthropic_error_event, diff --git a/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py b/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py index 59b8d40577a..e70d583a027 100644 --- a/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py +++ b/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py @@ -7,6 +7,7 @@ records the reported values without recomputing them. """ import math +from collections.abc import Mapping from dataclasses import dataclass from typing import Final @@ -89,7 +90,7 @@ def _parse_total_tokens(raw_value: str | None) -> int | None: return total_tokens -def parse_upstream_reported_usage(headers: httpx.Headers) -> UpstreamReportedUsage | None: +def parse_upstream_reported_usage(headers: Mapping[str, str]) -> UpstreamReportedUsage | None: """Read the reported totals off an upstream pass-through response. Returns ``None`` when neither header is present, which is the normal case diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index c751091ff9b..15e7b7e8482 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -5,10 +5,10 @@ Attachments define WHERE policies apply, separate from the policy definitions. This allows the same policy to be attached to multiple scopes. """ -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict +from typing import TYPE_CHECKING, Final, TypedDict from litellm._logging import verbose_proxy_logger from litellm.repositories.table_repositories import PolicyAttachmentRepository @@ -102,7 +102,7 @@ class AttachmentRegistry: self._initialized = True verbose_proxy_logger.info("Loaded %s policy attachments", len(self._attachments)) - def _parse_attachment(self, attachment_data: dict[str, Any]) -> PolicyAttachment: + def _parse_attachment(self, attachment_data: Mapping[str, object]) -> PolicyAttachment: """ Parse an attachment from raw configuration data. @@ -112,15 +112,17 @@ class AttachmentRegistry: Returns: Parsed PolicyAttachment object """ - return PolicyAttachment( - policy=attachment_data.get("policy", ""), - scope=attachment_data.get("scope"), - teams=attachment_data.get("teams"), - keys=attachment_data.get("keys"), - models=attachment_data.get("models"), - tags=attachment_data.get("tags"), - priority=attachment_data.get("priority"), - default=attachment_data.get("default", False), + return PolicyAttachment.model_validate( + { + "policy": attachment_data.get("policy", ""), + "scope": attachment_data.get("scope"), + "teams": attachment_data.get("teams"), + "keys": attachment_data.get("keys"), + "models": attachment_data.get("models"), + "tags": attachment_data.get("tags"), + "priority": attachment_data.get("priority"), + "default": attachment_data.get("default", False), + } ) def get_attached_policies( diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 1b1b2302236..4ad29bd3478 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -11,7 +11,7 @@ from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeVar, cast -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -52,6 +52,8 @@ try: except ImportError: HTTPException = None +_LEGACY_HOOK_REPLACEMENT: Final = TypeAdapter(object) + class UndeliverableStreamRewrite(Exception): def __init__(self, guardrail_name: str, reason: str) -> None: @@ -232,10 +234,12 @@ class _LegacyHookStreamAdapter(CustomGuardrail): logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response")) - replacement: Final = await self.inner.async_post_call_success_hook( - data=request_data, - user_api_key_dict=self.user_api_key_dict, - response=hooked, + replacement: Final = _LEGACY_HOOK_REPLACEMENT.validate_python( + await self.inner.async_post_call_success_hook( + data=request_data, + user_api_key_dict=self.user_api_key_dict, + response=hooked, + ) ) rewrite: Final = hooked if replacement is None else replacement if rewrite is None: diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index f4b38bea14e..68ac71c567b 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -623,7 +623,7 @@ async def test_pipeline( ``` """ try: - validated_pipeline: Final = GuardrailPipeline(**request.pipeline) + validated_pipeline: Final = GuardrailPipeline.model_validate(request.pipeline) except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid pipeline: {e}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3fe36f03207..5f5a53b3680 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -24,6 +24,7 @@ from collections.abc import ( Awaitable, Callable, Collection, + Container, Iterable, Mapping, MutableMapping, @@ -360,6 +361,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_parent_otel_span_from_kwargs, iter_items, normalize_drop_params, + require_attributes, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -474,6 +476,7 @@ from litellm.proxy.common_utils.codex_model_catalog import codex_model_list_body from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.credential_hydration import decrypted_or_stored +from litellm.proxy.common_utils.custom_openapi_spec import JsonObject from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.discoverable_model_filter import ( @@ -976,6 +979,7 @@ from litellm.secret_managers.main import ( ) from litellm.tracing.config import is_lens_tracing_enabled from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.anthropic import ( AnthropicMessagesRequest, AnthropicResponse, @@ -1139,7 +1143,7 @@ def _redact_worker_config_for_logging(worker_config: str | dict[str, JsonValue] return None if isinstance(worker_config, dict): return _redact_secret_values_in_obj(worker_config) - parsed: Final = safe_json_loads(worker_config, default=None) + parsed: Final[JsonValue] = safe_json_loads(worker_config, default=None) if isinstance(parsed, dict): return safe_dumps(_redact_secret_values_in_obj(parsed)) return worker_config @@ -1465,11 +1469,13 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ) _module_path, _func_name = _hook_spec.rsplit(":", 1) _module = importlib.import_module(_module_path) - _hook_fn = getattr(_module, _func_name) + _hook_fn = _UNTYPED_VALUE.validate_python(getattr(_module, _func_name)) if inspect.iscoroutinefunction(_hook_fn): await _hook_fn() - else: + elif callable(_hook_fn): _hook_fn() + else: + raise TypeError(f"'{type(_hook_fn).__name__}' object is not callable") verbose_proxy_logger.info("Worker startup hook '%s' executed successfully", _hook_spec) except Exception as e: verbose_proxy_logger.error("Worker startup hook '%s' failed: %s", _hook_spec, e) @@ -1917,15 +1923,22 @@ def _strip_operation_id_method_suffix(operation_id: str) -> str: return operation_id +class _SupportsValues(Protocol): + def values(self) -> Iterable[JsonValue]: ... + + +def _openapi_path_items(openapi_schema: JsonObject) -> Iterable[JsonObject]: + paths: Final = require_attributes(openapi_schema.get("paths", {}), _SupportsValues, "values") + return (path_item for path_item in paths.values() if isinstance(path_item, dict)) + + def ensure_unique_openapi_operation_ids( - openapi_schema: dict[str, Any], + openapi_schema: JsonObject, reserved_operation_ids: set[str] | None = None, -) -> dict[str, Any]: +) -> JsonObject: operation_entries: Final = [] operation_id_counts: Final[dict[str, int]] = {} - for path_item in openapi_schema.get("paths", {}).values(): - if not isinstance(path_item, dict): - continue + for path_item in _openapi_path_items(openapi_schema): for method, operation in path_item.items(): if method not in _OPENAPI_HTTP_METHODS or not isinstance(operation, dict): continue @@ -2002,6 +2015,19 @@ from starlette.routing import BaseRoute from litellm.proxy.common_utils.openapi_route import OpenAPISchema, serve_openapi_off_the_event_loop +class _OpenAPIQueryParam(Protocol): + @property + def name(self) -> object: ... + + @property + def required(self) -> object: ... + + +def _query_param_name_and_required(param: object) -> tuple[object, object]: + fields: Final = require_attributes(param, _OpenAPIQueryParam, "name", "required") + return fields.name, fields.required + + def _inject_websocket_stubs_into_openapi_schema(openapi_schema: dict, websocket_routes: list) -> dict: """ Add a synthetic GET stub for each WebSocket route so it appears in Swagger UI. @@ -2018,14 +2044,15 @@ def _inject_websocket_stubs_into_openapi_schema(openapi_schema: dict, websocket_ try: if hasattr(route, "dependant") and route.dependant is not None: # Handle both FastAPI <0.120 and >=0.120 - query_params = getattr(route.dependant, "query_params", []) + query_params = _UNTYPED_VALUE.validate_python(getattr(route.dependant, "query_params", [])) if query_params: - for param in query_params: + for param in iter_items(query_params): + name, required = _query_param_name_and_required(param) parameters.append( { - "name": param.name, + "name": name, "in": "query", - "required": param.required, + "required": required, "schema": {"type": "string"}, } ) @@ -5575,6 +5602,13 @@ def _log_ignored_cost_map_copy(model_id: str, fields: tuple[str, ...]) -> None: _DB_CONFIG_CALLBACK_EVENT_TYPES: Final[Mapping[str, tuple[Literal["success", "failure"], ...]]] = MappingProxyType( {"success_callback": ("success",), "failure_callback": ("failure",), "callbacks": ("success", "failure")} ) +_DB_CONFIG_CALLBACK_LISTS: Final[Mapping[str, Callable[[], Container[object]]]] = MappingProxyType( + { + "success_callback": lambda: litellm.success_callback, + "failure_callback": lambda: litellm.failure_callback, + "callbacks": lambda: litellm.callbacks, + } +) _CALLBACK_LIST_NAMES: Final = ( "callbacks", "success_callback", @@ -5585,7 +5619,9 @@ _CALLBACK_LIST_NAMES: Final = ( def _callback_list_entries(list_name: str) -> tuple[tuple[str, object], ...]: - return tuple((list_name, entry) for entry in getattr(litellm, list_name)) + return tuple( + (list_name, entry) for entry in iter_items(_UNTYPED_VALUE.validate_python(getattr(litellm, list_name))) + ) def _registered_callback_entries() -> tuple[tuple[str, object], ...]: @@ -5614,7 +5650,7 @@ def _restore_displaced_callback(list_name: str, entry: object) -> None: def _unregister_db_callback(registration: _DbCallbackRegistration) -> None: for list_name, entry in registration.added: litellm.logging_callback_manager.remove_callback_from_list_by_object( - getattr(litellm, list_name), entry, require_self=False + _UNTYPED_VALUE.validate_python(getattr(litellm, list_name)), entry, require_self=False ) for list_name, entry in registration.displaced: _restore_displaced_callback(list_name, entry) @@ -7823,7 +7859,7 @@ class ProxyConfig: self, callback: str, event_types: list[Literal["success", "failure"]], - existing_callbacks: list, + existing_callbacks: Container[object], ) -> None: """ Helper method to add a single callback to litellm for specified event types. @@ -7866,7 +7902,7 @@ class ProxyConfig: self._add_callback_from_db_to_in_memory_litellm_callbacks( callback=callback, event_types=list(_DB_CONFIG_CALLBACK_EVENT_TYPES[setting_key]), - existing_callbacks=getattr(litellm, setting_key), + existing_callbacks=_DB_CONFIG_CALLBACK_LISTS[setting_key](), ) after: Final = _registered_callback_entries() return _DbCallbackRegistration( @@ -8361,8 +8397,12 @@ class ProxyConfig: 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)) + response: _ConfigParamValueRow | None for section, response in zip(sections, responses): - if response is None or (param_value := getattr(response, "param_value", None)) is None: + if response is None: + continue + param_value: JsonValue = getattr(response, "param_value", None) + if param_value is None: continue verbose_proxy_logger.debug( @@ -8389,7 +8429,10 @@ class ProxyConfig: 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: + if current is None: + return + param_value: Final[object] = getattr(current, "param_value", None) + if param_value is None: return self._apply_db_environment_variables(param_value, _config_row_revision(current), environment_before_db, read) @@ -10383,6 +10426,7 @@ async def async_data_generator( request: Request | None = None, *, responses_stream_errors: bool = False, + dispatched_call_type: str | None = None, ): verbose_proxy_logger.debug("inside generator") stream_completed = False @@ -10598,6 +10642,7 @@ async def async_data_generator( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=request_data, + dispatched_call_type=dispatched_call_type, ) client_error: Final[Exception] = e if replacement is None else replacement verbose_proxy_logger.debug( @@ -11251,6 +11296,19 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + ### TOOL DISCOVERY (own lock, same interval, never rides the spend commit) ### + from litellm.proxy.utils import commit_tool_discoveries + + scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + commit_tool_discoveries, + "interval", + seconds=batch_writing_interval, + args=[prisma_client, proxy_logging_obj], + id="commit_tool_discoveries_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + ### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ### ## Reduces QPS as there are more tags for a single request tag_spend_update_interval: Final = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER) @@ -12207,7 +12265,7 @@ async def _deployment_hidden_by_listing_callbacks(deployment: Deployment, user_a async def model_list( request: Request = None, # pyright: ignore[reportArgumentType] # FastAPI always injects the Request; the None default only serves direct in-process callers user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - return_wildcard_routes: bool | None = False, + return_wildcard_routes: bool | None = None, team_id: str | None = None, include_model_access_groups: bool | None = False, only_model_access_groups: bool | None = False, @@ -12223,6 +12281,9 @@ async def model_list( This is just for compatibility with openai projects like aider. Query Parameters: + - return_wildcard_routes: When true, also list wildcard routes such as `openai/*`. + Defaults to `general_settings.return_wildcard_routes`, and an explicit + value on the request wins over it - client_version: Sent by Codex CLI (`?client_version=0.159.3`) when it fetches a provider's model catalog. When present, the response is Codex's own catalog shape (`{"models": [...]}`) built from the same listing, with @@ -12251,6 +12312,9 @@ async def model_list( global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj settings: Final = cast(dict[str, object], general_settings) # any-ok: legacy settings + include_wildcard_routes: Final = ( + return_wildcard_routes if return_wildcard_routes is not None else settings.get("return_wildcard_routes") is True + ) from litellm.llms.anthropic.common_utils import ( create_anthropic_model_list_response, @@ -12340,7 +12404,7 @@ async def model_list( proxy_model_list=proxy_model_list, user_model=None, infer_model_from_keys=False, - return_wildcard_routes=return_wildcard_routes or False, + return_wildcard_routes=include_wildcard_routes, llm_router=llm_router, model_access_groups=model_access_groups, include_model_access_groups=include_model_access_groups or False, @@ -12408,7 +12472,7 @@ async def model_list( team_id=team_id, include_model_access_groups=include_model_access_groups or False, only_model_access_groups=only_model_access_groups or False, - return_wildcard_routes=return_wildcard_routes or False, + return_wildcard_routes=include_wildcard_routes, user_api_key_cache=user_api_key_cache, ) @@ -15988,7 +16052,23 @@ async def _get_caller_byok_team_scope( return key_team_scope | set(user_row.teams or []) -def _byok_row_outside_caller_teams(model_info_dict: dict[str, JsonValue], allowed_team_ids: set[str] | None) -> bool: +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _with_get(value: object) -> _SupportsGet: + return require_attributes(value, _SupportsGet, "get") + + +class _SupportsLower(Protocol): + def lower(self) -> object: ... + + +def _lowered(value: object) -> object: + return require_attributes(value, _SupportsLower, "lower").lower() + + +def _byok_row_outside_caller_teams(model_info_dict: object, allowed_team_ids: set[str] | None) -> bool: """Whether a team BYOK row belongs to a team the caller is not a member of. `team_id` is only set on team BYOK rows; non-team rows fall through @@ -15996,7 +16076,7 @@ def _byok_row_outside_caller_teams(model_info_dict: dict[str, JsonValue], allowe """ if allowed_team_ids is None: return False - team_id: Final = model_info_dict.get("team_id") + team_id: Final = _with_get(model_info_dict).get("team_id") if team_id is None: return False return team_id not in allowed_team_ids @@ -16097,7 +16177,7 @@ async def _fetch_db_models_for_search( async def _apply_search_filter_to_models( - all_models: list[dict[str, Any]], + all_models: list[dict[str, object]], search: str, prisma_client: PrismaClient | None, proxy_config: ProxyConfig, @@ -16149,7 +16229,7 @@ async def _apply_search_filter_to_models( prisma_client=prisma_client, ) - def _is_byok_outside_caller_teams(model_info_dict: dict[str, JsonValue]) -> bool: + def _is_byok_outside_caller_teams(model_info_dict: object) -> bool: return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids) def _model_matches_search(m: dict[str, Any]) -> bool: @@ -16186,7 +16266,7 @@ async def _apply_search_filter_to_models( db_model_ids_in_router: Final = set() for m in filtered_router_models: - model_info = m.get("model_info", {}) + model_info = _with_get(m.get("model_info", {})) is_db_model = model_info.get("db_model", False) model_id = model_info.get("id") @@ -16297,7 +16377,7 @@ def _sort_models( reverse: Final = sort_order.lower() == "desc" def get_sort_key(model: dict[str, Any]) -> Any: - model_info: Final = model.get("model_info", {}) + model_info: Final = _with_get(_UNTYPED_VALUE.validate_python(model.get("model_info", {}))) if sort_by == "model_name": # Team BYOK models persist an internal `model_name` (e.g. @@ -16309,7 +16389,7 @@ def _sort_models( team_public_model_name: Final = model_info.get("team_public_model_name") if team_public_model_name: return str(team_public_model_name).lower() - return model.get("model_name", "").lower() + return _lowered(_UNTYPED_VALUE.validate_python(model.get("model_name", ""))) elif sort_by == "created_at": created_at: Final = model_info.get("created_at") @@ -16335,8 +16415,8 @@ def _sort_models( return normalized_dt elif sort_by == "costs": - input_cost: Final = model_info.get("input_cost_per_token", 0) or 0 - output_cost: Final = model_info.get("output_cost_per_token", 0) or 0 + input_cost: Final = model.get("model_info", {}).get("input_cost_per_token", 0) or 0 + output_cost: Final = model.get("model_info", {}).get("output_cost_per_token", 0) or 0 total_cost: Final = input_cost + output_cost # Put 0 or None costs at the end for asc, at the start for desc if total_cost == 0: @@ -16571,7 +16651,7 @@ async def _authorize_team_id_query( async def _filter_models_by_team_id( - all_models: list[dict[str, Any]], + all_models: list[dict[str, object]], team_id: str, prisma_client: PrismaClient, llm_router: Router, @@ -16615,7 +16695,7 @@ async def _filter_models_by_team_id( # every public model the admin can call. filtered_models: Final = [] for _model in all_models: - model_info = _model.get("model_info", {}) + model_info = _with_get(_model.get("model_info", {})) model_id = model_info.get("id", None) # BYOK rows owned by this team are always accessible to it, even if @@ -18219,6 +18299,7 @@ async def async_queue_request( try: data = await request.json() data.pop("_litellm_strip_stream_usage", None) + data.pop(STREAM_USAGE_INJECTED_FIELD, None) # Include original request and headers in the data data["proxy_server_request"] = { @@ -19653,6 +19734,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "proxy_config_reload_interval_seconds": "Integer", "pass_through_endpoints": "PydanticModel", "store_model_in_db": "Boolean", + "return_wildcard_routes": "Boolean", "store_prompts_in_spend_logs": "Boolean", "disable_requester_ip_address_logging": "Boolean", "disable_fallbacks_on_per_model_rate_limits": "Boolean", @@ -21713,8 +21795,7 @@ async def toolset_mcp_route(toolset_name: str, request: Request): detail=f"Toolset '{toolset_name}' not found", ) - scope: Final = dict(request.scope) - scope["path"] = "/mcp" + scope: Final = {**request.scope, "_original_path": request.scope.get("path", ""), "path": "/mcp"} token: Final = mcp_active_toolset_id.set(toolset.toolset_id) try: diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 4bec42a7868..19ee465c7dc 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1995,6 +1995,24 @@ ], "default_model_placeholder": "laya/english" }, + { + "provider": "OPENCODE_GO", + "provider_display_name": "OpenCode Go", + "litellm_provider": "opencode_go", + "credential_fields": [ + { + "key": "api_key", + "label": "API Key", + "placeholder": "OpenCode Console API key", + "tooltip": "Use a Console API key with an active Go or Go Plus subscription.", + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "glm-5.3-flash" + }, { "provider": "LAMBDA_AI", "provider_display_name": "Lambda Ai", diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 6ced7827846..351e4c10d87 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.get_blog_posts import ( get_blog_posts, ) from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.proxy._types import ( CommonProxyErrors, ) @@ -564,7 +565,7 @@ async def _resolve_autorouter_presets( url: str, fetch: Callable[[str], Awaitable[Mapping[str, AutoRouterPresetRecord]]], ) -> Mapping[str, AutoRouterPresetRecord]: - if os.getenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "").lower() == "true": + if is_offline_mode() or os.getenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "").lower() == "true": return _load_bundled_autorouter_presets() try: return await fetch(url) diff --git a/litellm/proxy/public_endpoints/public_v1/model_hub.py b/litellm/proxy/public_endpoints/public_v1/model_hub.py index cd634e490b6..e58149fba0b 100644 --- a/litellm/proxy/public_endpoints/public_v1/model_hub.py +++ b/litellm/proxy/public_endpoints/public_v1/model_hub.py @@ -107,10 +107,14 @@ def _features(row: ModelGroupInfoProxy) -> tuple[str, ...]: The hub's feature control has always been a multi-select over the `supports_*` flags. One boolean filter per flag would AND them, which is the opposite of what it does. """ + return _enabled_features(row.model_dump()) + + +def _enabled_features(fields: Mapping[str, object]) -> tuple[str, ...]: return tuple( sorted( name.removeprefix(FEATURE_PREFIX) - for name, value in row.model_dump().items() + for name, value in fields.items() if name.startswith(FEATURE_PREFIX) and value is True ) ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index c918ea3f1d1..312c5124ca5 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -37,6 +37,7 @@ from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADAT from litellm.integrations.custom_guardrail import ModifyResponseException 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 Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, first_validation_error from litellm.llms.openai.responses.guardrail_translation.handler import ( @@ -79,7 +80,8 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio 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.proxy.utils import count_request_input_tokens, handle_exception_on_proxy +from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task from litellm.router import Router @@ -94,7 +96,6 @@ from litellm.types.token_count import TokenCountSourceFields from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, TokenCountResponse if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.utils import ProxyLogging router: Final = APIRouter() @@ -532,6 +533,14 @@ async def cursor_model_list( return await model_list(user_api_key_dict=user_api_key_dict) +def _dispatched_prompt_token_counter(logging_obj: LiteLLMLoggingObj) -> Callable[[], int]: + return lambda: count_request_input_tokens( + model=str(logging_obj.model_call_details.get("model") or ""), + request_input=messages_since_newest_compaction(logging_obj.model_call_details.get("messages")), + system_input=None, + ) + + @router.post( "/cursor/chat/completions", dependencies=[ @@ -656,6 +665,11 @@ async def cursor_chat_completions( model=request_data.get("model", ""), custom_llm_provider=None, logging_obj=logging_obj, + count_prompt_tokens=( + _dispatched_prompt_token_counter(logging_obj) + if isinstance(logging_obj, LiteLLMLoggingObj) + else None + ), ) # Use async_data_generator to format as SSE return async_data_generator( @@ -663,6 +677,7 @@ async def cursor_chat_completions( user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + dispatched_call_type=processor.route_type, ) # Otherwise, use the default generator return async_data_generator( @@ -670,6 +685,7 @@ async def cursor_chat_completions( user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + dispatched_call_type=processor.route_type, ) try: @@ -1585,6 +1601,7 @@ def _responses_input_as_token_count_messages( transformed: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=input_value, responses_api_request=request_params, + only_since_compaction=True, ) return tuple( message if isinstance(message, dict) else message.model_dump(exclude_none=True) for message in transformed diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 787f27f56c8..e50587b1a0e 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -308,10 +308,7 @@ def should_use_polling_for_request( # Otherwise, check ALL deployments for this model_name in router elif llm_router is not None: try: - # Get all deployment indices for this model name - indices: Final = llm_router.model_name_to_deployment_indices.get(model, []) - for idx in indices: - deployment_dict = llm_router.model_list[idx] + for deployment_dict in llm_router.deployments_named(model) or (): litellm_params = deployment_dict.get("litellm_params", {}) # Check custom_llm_provider first diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 25f9f1dff15..cb226929167 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,8 +1,13 @@ import asyncio -from collections.abc import Mapping +import hashlib +import json +from collections import OrderedDict +from collections.abc import Callable, Mapping from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, Literal +from uuid import uuid4 import httpx from fastapi import HTTPException, status @@ -10,6 +15,7 @@ from pydantic import BeforeValidator, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import is_proxy_admin_request @@ -54,7 +60,34 @@ _ROUTER_ONLY_SETTINGS: Final = frozenset( ) -def _route_user_config_request(data: dict, route_type: str): +class UserConfigRouters: + def __init__(self, max_size: int) -> None: + self._max_size: Final = max_size + self._routers: Final[OrderedDict[str, LitellmRouter]] = OrderedDict() # mutable-ok: LRU evicts in place + + def get(self, config: Mapping[str, object], build: Callable[[], LitellmRouter]) -> LitellmRouter: + key: Final = hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest() + cached: Final = self._routers.get(key) + if cached is not None: + self._routers.move_to_end(key) + return cached + router: Final = build() + router.serve_through_request_logging() + self._routers[key] = router + if len(self._routers) > self._max_size: + self._routers.popitem(last=False)[1].discard() + return router + + def clear(self) -> None: + for router in self._routers.values(): + router.discard() + self._routers.clear() + + +user_config_routers: Final = UserConfigRouters(max_size=32) + + +async def _route_user_config_request(data: dict, route_type: str): """Route a request using the user-provided router config.""" router_config: Final = data.pop("user_config") @@ -63,10 +96,24 @@ def _route_user_config_request(data: dict, route_type: str): valid_args: Final = litellm.Router.get_valid_args() filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args} - user_router: Final = litellm.Router(**filtered_config) - ret_val: Final = getattr(user_router, f"{route_type}")(**data) - user_router.discard() - return ret_val + logging_obj: Final = _request_logging_obj(JSON_OBJECT_ADAPTER.validate_python(data), route_type) + cached_router: Final = user_config_routers.get(filtered_config, lambda: litellm.Router(**filtered_config)) + cached_router.bind_to_request_logging(logging_obj) + return await getattr(cached_router, f"{route_type}")(**{**data, "litellm_logging_obj": logging_obj}) + + +def _request_logging_obj(data: Mapping[str, object], route_type: str) -> LiteLLMLoggingObj: + arrived_with: Final = data.get("litellm_logging_obj") + if isinstance(arrived_with, LiteLLMLoggingObj): + return arrived_with + logging_obj, _ = litellm.utils.function_setup( + original_function=route_type, + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), # noqa: DTZ005 # Logging subtracts it from naive end times + is_async_call=True, + **{"litellm_call_id": str(uuid4()), **data}, + ) + return logging_obj def _is_a2a_agent_model(model_name: object) -> bool: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 803df9f83ca..0b91e57f01e 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -728,7 +728,9 @@ async def _get_team_member_budget_counter( membership_cache_key: Final = team_membership_reservation_cache_key( user_id=valid_token.user_id, team_id=team_object.team_id ) - cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key) + cached_team_membership: Final = _CACHED_VALUE.validate_python( + await user_api_key_cache.async_get_cache(key=membership_cache_key) + ) team_membership: LiteLLM_TeamMembership | None = None if isinstance(cached_team_membership, LiteLLM_TeamMembership): team_membership = cached_team_membership diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 6c7083fb6db..927152296e9 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -19,6 +19,10 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger +from litellm.responses.litellm_completion_transformation.compaction_items import ( + billable_input_since_newest_compaction, + holds_compaction_summary, +) from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.token_counter import ( @@ -54,10 +58,11 @@ async def count_input_tokens( (model, rust_tokenizer(model)) for model in models ) groups: Final[tuple[RustTokenizer | None, ...]] = tuple(dict.fromkeys(tokenizer for _, tokenizer in tokenizers)) + counted_raw_body: Final = None if holds_compaction_summary(request_body.get("input")) else raw_body group_counts: Final = [ await _count_group( request_body=request_body, - raw_body=raw_body, + raw_body=counted_raw_body, tokenizer=tokenizer, models=tuple(model for model, selected in tokenizers if selected == tokenizer), ) @@ -140,7 +145,9 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: - return _count_text_tokens(model=model, text=request_body.get("input")) + return _count_text_tokens( + model=model, text=billable_input_since_newest_compaction(request_body.get("input")) + ) if "query" in request_body or "documents" in request_body: query_tokens: Final = _count_text_tokens(model=model, text=request_body.get("query")) document_tokens: Final = _count_text_tokens( diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index ea543aeae6e..4abc79febe1 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -21,6 +21,8 @@ from itertools import chain, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol +from pydantic import TypeAdapter + from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.constants import ( @@ -135,6 +137,9 @@ def _public_model_name(row: object, model_info: Mapping[str, object]) -> str: return str(getattr(row, "model_name", "") or "") +_DECODED_JSON: Final = TypeAdapter(object) + + def _decode_model_info(raw: object) -> "Mapping[str, object] | None": """A deployment's model_info as a mapping, decoding a JSON string, else None. @@ -143,7 +148,7 @@ def _decode_model_info(raw: object) -> "Mapping[str, object] | None": """ if isinstance(raw, str): try: - decoded: Final[object] = json.loads(raw) + decoded: Final = _DECODED_JSON.validate_python(json.loads(raw)) except (TypeError, ValueError): return None return decoded if isinstance(decoded, dict) else None diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 2aab8deb034..604a7a389e6 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -19,6 +19,7 @@ from pydantic import ConfigDict, Field import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_prompt_caching_savings, generic_cost_per_token, @@ -69,7 +70,9 @@ def _resolve_model(model: str | None, custom_llm_provider: str | None) -> _Model if not model: return None try: - resolved_model, provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + resolved_model, provider = resolve_model_and_provider_for_metadata( + model=model, custom_llm_provider=custom_llm_provider + ) except Exception as e: # noqa: BLE001 # get_llm_provider raises for unroutable names; degrade to an unavailable estimate verbose_proxy_logger.debug( "savings: cannot resolve provider for model=%s custom_llm_provider=%s (%s)", model, custom_llm_provider, e diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 84278f4619d..3c65fc0bf06 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3045,7 +3045,7 @@ async def ui_view_spend_logs( ) sql_params.extend([page_size, skip]) - data: Final = await prisma_client.db.query_raw(sql_query, *sql_params) + data: Final[Sequence[Mapping[str, object]]] = await _query_raw(prisma_client, sql_query, *sql_params) if request_id is not None and not is_v2 and not is_admin_view: await _assert_user_owns_fetched_spend_rows( @@ -3064,7 +3064,7 @@ async def ui_view_spend_logs( return await _build_ui_spend_logs_response( prisma_client, - data, + list(data), total_records, page, page_size, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 62aee788ed0..94b079163e9 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -27,6 +27,7 @@ from litellm.constants import ( LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, + MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB, MCP_SPEND_LOG_MODEL_PREFIX, NON_INFERENCE_CALL_TYPES, REDACTED_BY_LITELM_STRING, @@ -70,6 +71,11 @@ from litellm.proxy._types import ( from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + summary_only_compaction_content, + with_compaction_summary, +) from litellm.types.passthrough_endpoints.pass_through_endpoints import ( SETTLES_BACKGROUND_RESPONSE_KWARG, BackgroundResponseSettlement, @@ -1227,11 +1233,42 @@ def _is_spend_log_content(parent: Mapping[str, object], key: str) -> bool: ) +def _truncated_for_db(value: str, max_length: int) -> str: + if len(value) <= max_length: + return value + start_chars: Final = int(max_length * 0.35) + end_chars: Final = int(max_length * 0.65) + return ( + f"{value[:start_chars]}" + f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {len(value) - start_chars - end_chars} chars. " + f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." + f"{value[-min(end_chars, max_length - start_chars) :]}" + ) + + +def _compaction_summary_for_db(summary: str) -> str: + return _truncated_for_db( + summary, max(_get_max_string_length_prompt_in_db(), MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB) + ) + + +def _bounded_compaction_summary( + parent: Mapping[str, object], key: str, value: object, keep_compaction_blocks: bool +) -> str | None: + if not keep_compaction_blocks or not isinstance(value, str) or parent.get("type") != "compaction": + return None + if key == "content": + return _compaction_summary_for_db(value) + summary: Final = compaction_summary(parent) if key == "encrypted_content" else None + return None if summary is None else with_compaction_summary(value, _compaction_summary_for_db(summary)) + + def _sanitize_request_body_for_spend_logs_payload( request_body: Mapping[str, object], visited: set | None = None, max_string_length_prompt_in_db: int | None = None, mask_credentials: bool = True, + keep_compaction_blocks: bool = False, ) -> dict: """ Recursively sanitize request body to prevent logging large base64 strings or other large values. @@ -1242,11 +1279,6 @@ def _sanitize_request_body_for_spend_logs_payload( SensitiveDataMasker classifies as credentials with REDACTED_BY_LITELM_STRING, except inside tool payloads and logprobs. """ - from litellm.constants import ( - LITELLM_TRUNCATED_PAYLOAD_FIELD, - LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, - ) - if visited is None: visited = set() if max_string_length_prompt_in_db is None: @@ -1261,51 +1293,23 @@ def _sanitize_request_body_for_spend_logs_payload( def _sanitize_value(value: object, mask_credentials: bool) -> object: if isinstance(value, Mapping): return _sanitize_request_body_for_spend_logs_payload( - value, visited, max_string_length_prompt_in_db, mask_credentials + value, visited, max_string_length_prompt_in_db, mask_credentials, keep_compaction_blocks ) elif isinstance(value, list): return [_sanitize_value(item, mask_credentials) for item in value] elif isinstance(value, str): - if len(value) > max_string_length_prompt_in_db: - # Keep 35% from beginning and 65% from end (end is usually more important) - # This split ensures we keep more context from the end of conversations - start_ratio: Final = 0.35 - end_ratio: Final = 0.65 - - # Calculate character distribution - start_chars: Final = int(max_string_length_prompt_in_db * start_ratio) - end_chars = int(max_string_length_prompt_in_db * end_ratio) - - # Ensure we don't exceed the total limit - total_keep: Final = start_chars + end_chars - if total_keep > max_string_length_prompt_in_db: - end_chars = max_string_length_prompt_in_db - start_chars - - # If the string length is less than what we want to keep, just truncate normally - if len(value) <= max_string_length_prompt_in_db: - return value - - # Calculate how many characters are being skipped - skipped_chars: Final = len(value) - total_keep - - # Build the truncated string: beginning + truncation marker + end - truncated_value: Final = ( - f"{value[:start_chars]}" - f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. " - f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." - f"{value[-end_chars:]}" - ) - return truncated_value - return value + return _truncated_for_db(value, max_string_length_prompt_in_db) return value - return { - k: REDACTED_BY_LITELM_STRING - if mask_credentials and _is_request_body_credential(k, v) - else _sanitize_value(v, mask_credentials and not _is_spend_log_content(request_body, k)) - for k, v in request_body.items() - if k not in _SENSITIVE_REQUEST_BODY_KEYS - } + def _stored_value(key: str, value: object) -> object: + if mask_credentials and _is_request_body_credential(key, value): + return REDACTED_BY_LITELM_STRING + bounded_summary: Final = _bounded_compaction_summary(request_body, key, value, keep_compaction_blocks) + if bounded_summary is not None: + return bounded_summary + return _sanitize_value(value, mask_credentials and not _is_spend_log_content(request_body, key)) + + return {k: _stored_value(k, v) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS} # Quoted-key form: ``"input"`` / ``'messages'`` / ``"prompt"`` followed by @@ -1820,6 +1824,33 @@ def _stored_request_without_requester_ip_address(request_body: Mapping[object, o } +def _with_bounded_compaction_summary(item: object, stored_item: object) -> object: + summary: Final = compaction_summary(item) + if summary is None or not isinstance(stored_item, Mapping): + return stored_item + return { + **_OBJECT_MAPPING.validate_python(stored_item), + "encrypted_content": summary_only_compaction_content(_compaction_summary_for_db(summary)), + } + + +def _sanitize_proxy_server_request_body(request_body: Mapping[str, object]) -> Mapping[object, object]: + sanitized: Final = _OBJECT_MAPPING.validate_python(_sanitize_request_body_for_spend_logs_payload(request_body)) + items: Final = request_body.get("input") + stored_items: Final = sanitized.get("input") + if not isinstance(items, list) or not isinstance(stored_items, list): + return sanitized + return { + **sanitized, + "input": [ + _with_bounded_compaction_summary(item, stored_item) + for item, stored_item in zip( + _OBJECT_LIST.validate_python(items), _OBJECT_LIST.validate_python(stored_items) + ) + ], + } + + def _get_proxy_server_request_for_spend_logs_payload( metadata: dict, litellm_params: Mapping[str, object], @@ -1864,8 +1895,7 @@ def _get_proxy_server_request_for_spend_logs_payload( _request_body = _convert_mapping_to_json_serializable(without_classifier_audit(_request_body)) perform_redaction(model_call_details=_request_body, result=None) - _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) - stored_request_body: Final = _OBJECT_MAPPING.validate_python(_request_body) + stored_request_body: Final = _sanitize_proxy_server_request_body(_request_body) _request_body_json_str: Final = safe_dumps( _stored_request_without_requester_ip_address(stored_request_body) if drop_requester_ip_address @@ -1939,7 +1969,9 @@ def _get_response_for_spend_logs_payload( response_obj = _convert_to_json_serializable_dict(response_obj) response_obj = perform_redaction(model_call_details={}, result=response_obj) - sanitized_wrapper: Final = _sanitize_request_body_for_spend_logs_payload({"response": response_obj}) + sanitized_wrapper: Final = _sanitize_request_body_for_spend_logs_payload( + {"response": response_obj}, keep_compaction_blocks=True + ) sanitized_response: Final = sanitized_wrapper.get("response", response_obj) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index d3858f95243..6938588444c 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -122,7 +122,9 @@ async def _get_vantage_settings(): encrypted_api_key: Final = settings.get("api_key") if encrypted_api_key: - decrypted_api_key = decrypt_value_helper(encrypted_api_key, key="vantage_api_key", exception_type="error") + decrypted_api_key: Final = decrypt_value_helper( + encrypted_api_key, key="vantage_api_key", exception_type="error" + ) if decrypted_api_key is None: raise HTTPException( status_code=500, diff --git a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py index 7411f1b4947..4bb17fd290e 100644 --- a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py @@ -1,5 +1,7 @@ import asyncio +import os import re +import xml.etree.ElementTree as ET from collections import Counter from collections.abc import Awaitable, Mapping from dataclasses import dataclass @@ -12,12 +14,14 @@ from pydantic import ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.llms.base import LiteLLMBaseModel router: Final = APIRouter() -LATEST_RELEASE_URL: Final = "https://api.github.com/repos/BerriAI/litellm/releases/latest" +CHANGELOG_URL: Final = "https://docs.litellm.ai/release_notes/" +CHANGELOG_FEED_URL: Final = f"{CHANGELOG_URL}rss.xml" RELEASE_BY_TAG_URL: Final = "https://api.github.com/repos/BerriAI/litellm/releases/tags/{tag}" LATEST_RELEASE_FETCH_TIMEOUT_SECONDS: Final = 5 LATEST_RELEASE_CACHE_TTL_SECONDS: Final = 60 * 60 @@ -26,7 +30,7 @@ LATEST_RELEASE_CACHE_KEY: Final = "latest_release_info" _RELEASE_BULLET_PATTERN: Final = re.compile(r"^\*\s+(?:([A-Za-z]+)(?:\([^)]*\))?!?:\s)?\S") _NEW_CONTRIBUTOR_PATTERN: Final = re.compile(r"^\*\s+@\S+ made their first contribution\b") -_STABLE_TAG_PATTERN: Final = re.compile(r"^v?(\d+)\.(\d+)\.(\d+)$") +_CHANGELOG_PAGE_PATTERN: Final = re.compile(rf"{re.escape(CHANGELOG_URL)}v(\d+)\.(\d+)\.(\d+)/") _Bucket: TypeAlias = Literal["new_features", "bug_fixes", "other_updates"] _PREFIX_BUCKETS: Final[Mapping[str, _Bucket]] = MappingProxyType({"feat": "new_features", "fix": "bug_fixes"}) @@ -45,6 +49,16 @@ class LatestReleaseUnavailable: reason: str +@dataclass(frozen=True, slots=True) +class ChangelogEntry: + version: tuple[int, int, int] + url: str + + @property + def tag(self) -> str: + return "v" + ".".join(str(part) for part in self.version) + + class _GitHubRelease(LiteLLMBaseModel): tag_name: str html_url: str @@ -76,6 +90,10 @@ def _default_fetch_lock() -> asyncio.Lock: return _latest_release_fetch_lock +def _default_release_check_enabled() -> bool: + return not is_offline_mode() and os.getenv("LITELLM_DISABLE_LATEST_RELEASE_CHECK", "").strip().lower() != "true" + + def _bucket_for(line: str) -> _Bucket | None: if _NEW_CONTRIBUTOR_PATTERN.match(line) is not None: return None @@ -91,9 +109,18 @@ def count_release_bullets(body: str) -> Mapping[_Bucket, int]: return MappingProxyType(Counter(bucket for line in body.splitlines() if (bucket := _bucket_for(line)) is not None)) -def _stable_version(tag_name: str) -> tuple[int, int, int] | None: - match: Final = _STABLE_TAG_PATTERN.match(tag_name) - return None if match is None else (int(match.group(1)), int(match.group(2)), int(match.group(3))) +def _changelog_entry(link: str) -> ChangelogEntry | None: + match: Final = _CHANGELOG_PAGE_PATTERN.match(link) + if match is None: + return None + return ChangelogEntry(version=(int(match.group(1)), int(match.group(2)), int(match.group(3))), url=link) + + +def latest_changelog_entry(feed: bytes) -> ChangelogEntry | None: + """The highest stable release with a page in the changelog feed. Release candidate pages don't count.""" + links: Final = (link.text or "" for link in ET.fromstring(feed).iterfind("channel/item/link")) + entries: Final = (entry for link in links if (entry := _changelog_entry(link)) is not None) + return max(entries, key=lambda entry: entry.version, default=None) def earlier_patch_tags(version: tuple[int, int, int]) -> tuple[str, ...]: @@ -102,7 +129,9 @@ def earlier_patch_tags(version: tuple[int, int, int]) -> tuple[str, ...]: return tuple(f"v{major}.{minor}.{earlier}" for earlier in range(patch)) -def summarize_release_line(latest: _GitHubRelease, earlier: tuple[_GitHubRelease, ...]) -> LatestReleaseInfo: +def summarize_release_line( + entry: ChangelogEntry, latest: _GitHubRelease, earlier: tuple[_GitHubRelease, ...] +) -> LatestReleaseInfo: """A patch release's notes only list its backports, so count the whole minor line it belongs to.""" counts: Final = sum( (Counter[_Bucket](count_release_bullets(release.body or "")) for release in (latest, *earlier)), @@ -113,7 +142,7 @@ def summarize_release_line(latest: _GitHubRelease, earlier: tuple[_GitHubRelease new_features=counts.get("new_features", 0), bug_fixes=counts.get("bug_fixes", 0), other_updates=counts.get("other_updates", 0), - release_url=latest.html_url, + release_url=entry.url, ) @@ -137,25 +166,38 @@ async def _fetch_earlier_release(client: _AsyncGetClient, tag: str) -> _GitHubRe return release +async def _fetch_changelog_entry(client: _AsyncGetClient) -> ChangelogEntry | LatestReleaseUnavailable: + response: Final = await client.get(CHANGELOG_FEED_URL, timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) + if response.status_code != 200: + return LatestReleaseUnavailable(reason=f"Changelog feed responded with status {response.status_code}") + try: + entry: Final = latest_changelog_entry(response.content) + except ET.ParseError as e: + return LatestReleaseUnavailable(reason=f"Changelog feed was not valid XML: {e}") + return LatestReleaseUnavailable(reason="Changelog feed lists no stable release") if entry is None else entry + + async def fetch_latest_release(client: _AsyncGetClient) -> LatestReleaseInfo | LatestReleaseUnavailable: try: + entry: Final = await _fetch_changelog_entry(client) + if isinstance(entry, LatestReleaseUnavailable): + return entry latest: Final = parse_github_release( - await client.get(LATEST_RELEASE_URL, timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) + await client.get(RELEASE_BY_TAG_URL.format(tag=entry.tag), timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) ) if isinstance(latest, LatestReleaseUnavailable): return latest - version: Final = _stable_version(latest.tag_name) - if version is None: - return summarize_release_line(latest, ()) earlier: Final = await asyncio.gather( - *(_fetch_earlier_release(client, tag) for tag in earlier_patch_tags(version)) + *(_fetch_earlier_release(client, tag) for tag in earlier_patch_tags(entry.version)) ) except httpx.HTTPError as e: return LatestReleaseUnavailable(reason=f"{type(e).__name__}: {e}") unavailable: Final = next((release for release in earlier if isinstance(release, LatestReleaseUnavailable)), None) if unavailable is not None: return unavailable - return summarize_release_line(latest, tuple(release for release in earlier if isinstance(release, _GitHubRelease))) + return summarize_release_line( + entry, latest, tuple(release for release in earlier if isinstance(release, _GitHubRelease)) + ) async def get_latest_release_info( @@ -188,12 +230,16 @@ async def latest_release_info( client: Annotated[_AsyncGetClient, Depends(_default_client)], cache: Annotated[InMemoryCache, Depends(_default_cache)], fetch_lock: Annotated[asyncio.Lock, Depends(_default_fetch_lock)], + release_check_enabled: Annotated[bool, Depends(_default_release_check_enabled)], ) -> LatestReleaseInfo | None: """ - Latest stable LiteLLM GitHub release with the PR count of its whole minor line (e.g. v1.105.0 + v1.105.1) - split into new features, bug fixes and other updates. - Returns null when GitHub can't be reached so the dashboard upgrade banner simply doesn't render. + Newest stable release on the LiteLLM changelog, linked to its changelog page, with the PR count of its whole + minor line (e.g. v1.105.0 + v1.105.1) from GitHub split into new features, bug fixes and other updates. + Returns null when the changelog or GitHub can't be reached so the dashboard upgrade banner simply doesn't render. + Set LITELLM_DISABLE_LATEST_RELEASE_CHECK=true to return null without making any outbound request. """ + if not release_check_enabled: + return None result: Final = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) if isinstance(result, LatestReleaseUnavailable): verbose_proxy_logger.warning("LiteLLM: latest release info unavailable: %s", result.reason) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index d51e1167888..e04ddbc75cc 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -8,6 +8,7 @@ from types import MappingProxyType from typing import ( Annotated, Final, + Literal, NamedTuple, Protocol, cast, # noqa: TID251 # prisma types Json columns as fields.Json but de-serializes them to plain python on read @@ -256,6 +257,11 @@ class UISettings(LiteLLMBaseModel): model_config = ConfigDict(extra="allow") + usage_default_time_range: Literal["today", "last_7_days", "last_30_days", "month_to_date", "year_to_date"] = Field( + default="last_7_days", + description="Default date range when opening Usage without an explicit URL range. Calendar ranges use the browser timezone.", + ) + disable_model_add_for_internal_users: bool = Field( default=False, description=( @@ -386,6 +392,7 @@ class UISettingsResponse(SettingsResponse): # Allowlist of UI settings that can be stored ALLOWED_UI_SETTINGS_FIELDS: Final = { + "usage_default_time_range", "disable_model_add_for_internal_users", "disable_team_admin_delete_team_user", "enabled_ui_pages_internal_users", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ea1044d4a55..ce7e6afde24 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -19,6 +19,7 @@ from collections.abc import ( Awaitable, Callable, Coroutine, + Iterable, Iterator, Mapping, Sequence, @@ -51,7 +52,7 @@ from typing import ( import anyio import httpx -from typing_extensions import Never, ReadOnly, TypedDict +from typing_extensions import Never, NotRequired, ReadOnly, TypedDict from litellm import _custom_logger_compatible_callbacks_literal from litellm.constants import ( @@ -89,6 +90,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( with_litellm_call_id, ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error +from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage @@ -158,7 +160,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int -from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -182,6 +184,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db import tool_registry_writer from litellm.proxy.db.create_views import ( create_missing_views, create_view_tolerating_race, @@ -189,6 +192,7 @@ from litellm.proxy.db.create_views import ( ) from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.db.db_url_settings import ( DatabaseURLSettings, add_missing_query_params, @@ -322,10 +326,22 @@ class _RelTuplesRow(TypedDict): reltuples: ReadOnly[int] +class _TeamMemberRow(TypedDict): + role: ReadOnly[Literal["admin", "user"]] + user_id: ReadOnly[NotRequired[str | None]] + user_email: ReadOnly[NotRequired[str | None]] + + +class _TeamMembershipRow(TypedDict): + user_id: ReadOnly[str | None] + team_members_with_roles: ReadOnly[Iterable[_TeamMemberRow] | None] + + _VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 _VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' _VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) +_DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] _ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception @@ -363,6 +379,47 @@ class _UpdateManyBatch(Protocol): async def commit(self) -> None: ... +def _update_many_batch(db: PrismaWrapper | RoutingPrismaWrapper) -> _UpdateManyBatch: + return db.batch_() + + +def _query_raw_rows( + db: PrismaWrapper | RoutingPrismaWrapper, query: str, *args: object +) -> Awaitable[Sequence[Mapping[str, object]]]: + return db.query_raw(query, *args) + + +def _query_first_row( + db: PrismaWrapper | RoutingPrismaWrapper, query: str, *args: object +) -> Awaitable[Mapping[str, object] | None]: + return db.query_first(query, *args) + + +def _team_member_for_user(row: _TeamMembershipRow) -> Member | None: + members_with_roles: Final = row["team_members_with_roles"] + if members_with_roles is None: + return None + user_id: Final = row["user_id"] + if user_id is None: + return None + matches: Final = [ + Member(**member) + for member in members_with_roles + if member.get("user_id") is not None and user_id == member.get("user_id") + ] + return matches[-1] if matches else None + + +class _EndUserSpendTransaction(Protocol): + def batch_(self) -> contextlib.AbstractAsyncContextManager[_EndUserSpendBatch]: ... + + +def _end_user_spend_transaction( + db: PrismaWrapper | RoutingPrismaWrapper, timeout: timedelta +) -> contextlib.AbstractAsyncContextManager[_EndUserSpendTransaction]: + return db.tx(timeout=timeout) + + unified_guardrail: Final = UnifiedLLMGuardrails() NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages}) @@ -542,7 +599,7 @@ class _UpstreamStreamBoundary(Generic[_T]): self.failure: BaseException | None = None def __getattr__(self, name: str) -> object: - return getattr(self._source, name) + return _DELEGATED_ATTRIBUTE.validate_python(getattr(self._source, name)) def __aiter__(self) -> "_UpstreamStreamBoundary[_T]": return self @@ -937,7 +994,7 @@ def _system_prompt_text(system_input: object) -> str: return "".join(_prompt_block_text(block) for block in system_input) -def _count_request_input_tokens(model: str, request_input: object, system_input: object) -> int: +def count_request_input_tokens(model: str, request_input: object, system_input: object) -> int: system_text: Final = _system_prompt_text(system_input) system_tokens: Final = litellm.token_counter(model=model, text=system_text) if system_text else 0 if isinstance(request_input, str): @@ -958,7 +1015,7 @@ def _estimate_dispatched_failure_usage(model: str, request_input: object, system the same tokenizer fallback interrupted streams use, so the spend log's failure row records what was sent instead of zero.""" try: - input_tokens: Final = _count_request_input_tokens( + input_tokens: Final = count_request_input_tokens( model=model, request_input=request_input, system_input=system_input ) except Exception: @@ -997,6 +1054,8 @@ _INPUT_ESTIMABLE_CALL_TYPES: Final = frozenset( ) +_RESPONSES_CALL_TYPES: Final = frozenset((CallTypes.responses.value, CallTypes.aresponses.value)) + _FAILED_ATTEMPTS_USAGE: Final = TypeAdapter(tuple[tuple[Usage, float], ...]) @@ -1015,13 +1074,14 @@ def _billed_across_attempts( ) -def _failure_usage_to_lift( +def _current_failure_usage_to_lift( model_call_details: Mapping[str, object], request_body: Mapping[str, object], dispatched: bool, provider_billed: bool = False, deployment_id: str | None = None, base_model: str | None = None, + dispatched_call_type: str | None = None, ) -> tuple[object, object] | None: """A stream that broke mid-flight still billed the provider for the chunks already delivered; the streaming handler stashes that recovered usage and @@ -1037,9 +1097,6 @@ def _failure_usage_to_lift( failure stays at zero cost. Returns the (combined_usage_object, response_cost) pair to lift, or None.""" recovered_usage: Final = model_call_details.get("combined_usage_object") - failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) - if failed_attempts: - return _billed_across_attempts(failed_attempts, recovered_usage, model_call_details.get("response_cost")) if recovered_usage is not None: return recovered_usage, model_call_details.get("response_cost") if not dispatched: @@ -1048,7 +1105,7 @@ def _failure_usage_to_lift( # The proxy-only error handler stamps this flag on a client disconnect too, even though the # request had already reached the provider, so a billed failure keeps its estimate. return None - estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body) + estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body, dispatched_call_type) if estimated_usage is None: return None if not provider_billed: @@ -1059,13 +1116,14 @@ def _failure_usage_to_lift( def _estimated_prompt_usage( - model_call_details: Mapping[str, object], request_body: Mapping[str, object] + model_call_details: Mapping[str, object], request_body: Mapping[str, object], dispatched_call_type: str | None ) -> Usage | None: """The input-only Usage estimate for the prompt this request sent, or None when the call type has no estimable input. The raw request body backfills the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a chat-completions provider) without filling optional_params.""" if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES: return None + dispatched_input: Final = model_call_details.get("messages") optional_params: Final = model_call_details.get("optional_params") dispatched_system: Final = ( (optional_params.get("system") or optional_params.get("instructions")) @@ -1082,7 +1140,11 @@ def _estimated_prompt_usage( ) return _estimate_dispatched_failure_usage( model=str(model_call_details.get("model") or ""), - request_input=model_call_details.get("messages"), + request_input=( + messages_since_newest_compaction(dispatched_input) + if dispatched_call_type in _RESPONSES_CALL_TYPES + else dispatched_input + ), system_input=system_input, ) @@ -1281,6 +1343,31 @@ def _estimated_input_cost( return next((cost for cost in candidate_costs if cost is not None), 0.0) +def _failure_usage_to_lift( + model_call_details: Mapping[str, object], + request_body: Mapping[str, object], + dispatched: bool, + provider_billed: bool = False, + deployment_id: str | None = None, + base_model: str | None = None, + dispatched_call_type: str | None = None, +) -> tuple[object, object] | None: + current_attempt: Final = _current_failure_usage_to_lift( + model_call_details, + request_body, + dispatched, + provider_billed, + deployment_id, + base_model, + dispatched_call_type=dispatched_call_type, + ) + failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) + if not failed_attempts: + return current_attempt + current_usage, current_cost = current_attempt if current_attempt is not None else (None, None) + return _billed_across_attempts(failed_attempts, current_usage, current_cost) + + _EMPTY_LIFT: Final = MappingProxyType({}) @@ -1389,6 +1476,13 @@ def _call_type_for_route(route: str | None) -> str | None: return call_types[0].value if len(operations) == 1 else None +REJECTED_BEFORE_REQUEST_SETUP_KEY: Final = "litellm_rejected_before_request_setup" +REJECTED_BEFORE_REQUEST_SETUP: Final = object() +"""Stored under ``REJECTED_BEFORE_REQUEST_SETUP_KEY`` on ``request_data`` when the failure handler finds no logging +object, meaning the request died (auth, body parsing, model resolution) before the pipeline set one up. A sentinel +object rather than ``True`` because ``request_data`` is the client body: JSON cannot forge this identity, so a +client can neither hide a guardrail block nor force a row by sending the key or a ``litellm_call_id`` itself.""" + _PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) _LOG_DB_METRICS_CALL_TYPES: Final = frozenset(("get_data", "insert_data", "update_data", "delete_data")) @@ -1424,7 +1518,9 @@ def _provider_billed_the_failure(exception: Exception) -> bool: return _timed_out_reading_from_provider(exception) or _client_disconnected(exception) -def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: bool = False) -> Mapping[str, object]: +def _failure_fields_to_lift( + request_data: Mapping[str, object], provider_billed: bool = False, dispatched_call_type: str | None = None +) -> Mapping[str, object]: """Failure-path callbacks run after ``litellm_logging_obj`` is popped from request_data (it is not serialisable), so the caller merges these fields onto request_data first: the request start and first-handoff instants for @@ -1444,13 +1540,14 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: _usage_to_lift: Final = _failure_usage_to_lift( model_call_details=_model_call_details, request_body=request_data, - dispatched=_first_handoff is not None, + dispatched=_model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None, provider_billed=_billed, deployment_id=_deployment_id, base_model=_base_model, + dispatched_call_type=dispatched_call_type, ) _prompt_estimate: Final = ( - _estimated_prompt_usage(_model_call_details, request_data) + _estimated_prompt_usage(_model_call_details, request_data, dispatched_call_type) if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") else None ) @@ -3614,6 +3711,7 @@ class ProxyLogging: error_type: ProxyErrorTypes | None = None, route: str | None = None, traceback_str: str | None = None, + dispatched_call_type: str | None = None, ) -> HTTPException | None: """ Allows users to raise custom exceptions/log when a call fails, without having to deal with parsing Request body. @@ -3684,7 +3782,9 @@ class ProxyLogging: request_data.update( await offload_token_count(_failure_fields_to_lift)( - request_data, provider_billed=_provider_billed_the_failure(original_exception) + request_data, + provider_billed=_provider_billed_the_failure(original_exception), + dispatched_call_type=dispatched_call_type or _call_type_for_route(user_api_key_dict.request_route), ) ) @@ -3795,6 +3895,7 @@ class ProxyLogging: **request_data, ) request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller + request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP}) if "metadata" not in request_data: request_data["metadata"] = {} request_data["metadata"].update(user_api_key_logged_metadata) @@ -4732,19 +4833,25 @@ class ProxyLogging: @staticmethod def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool: - """Drop the parked success dispatch for an assembled chat stream that ends in an error - ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. + """Drop the parked success dispatch for an assembled chat stream or a stream that billed several rounds when it + ends in an error ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. Returns False when the parked dispatch should still be flushed by the caller.""" logging_obj: Final = request_data.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): return False _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None) - assembled: Final = _args[0] if _args else None - if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse): + if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS): return False + match _args: + case (ModelResponse() as assembled, *_): + logging_obj.record_assembled_response_for_failure(assembled) + case (BilledRoundsDispatch() as billed,): + billed.dispatch.close() + logging_obj.record_billed_rounds_for_failure(billed.spend) + case _: + return False logging_obj._on_deferred_stream_complete = None logging_obj._deferred_stream_complete_args = None - logging_obj.record_assembled_response_for_failure(assembled) return True async def arelease_max_parallel_requests_on_disconnect( @@ -5451,7 +5558,9 @@ class PrismaClient: raise e - async def _query_first_with_cached_plan_fallback(self, sql_query: str, *args) -> dict | None: + async def _query_first_with_cached_plan_fallback( + self, sql_query: str, *args: object + ) -> Mapping[str, object] | None: """ Execute a query, recovering once from PostgreSQL's "cached plan must not change result type" error. @@ -5702,7 +5811,7 @@ class PrismaClient: LIMIT $1 OFFSET $2 """ - response = await self.db.query_raw(sql_query, limit, offset) + response = await _query_raw_rows(self.db, sql_query, limit, offset) return response elif table_name == "spend": verbose_proxy_logger.debug("PrismaClient: get_data: table_name == 'spend'") @@ -5846,7 +5955,7 @@ class PrismaClient: """ response = ( - await self.writer_db.query_first(sql_query, hashed_token) + await _query_first_row(self.writer_db, sql_query, hashed_token) if use_writer else await self._query_first_with_cached_plan_fallback(sql_query, hashed_token) ) @@ -5882,31 +5991,14 @@ class PrismaClient: if response["team_blocked"] is None: response["team_blocked"] = False - team_member: Member | None = None - if response["team_members_with_roles"] is not None and response["user_id"] is not None: - ## find the team member corresponding to user id - """ - [ - { - "role": "admin", - "user_id": "default_user_id", - "user_email": null - }, - { - "role": "user", - "user_id": null, - "user_email": "test@email.com" - } - ] - """ - for tm in response["team_members_with_roles"]: - if tm.get("user_id") is not None and response["user_id"] == tm.get("user_id"): - team_member = Member(**tm) + ## find the team member corresponding to user id + team_member: Final = _team_member_for_user(response) response["team_member"] = team_member - response = LiteLLM_VerificationTokenView(**response, last_refreshed_at=time.time()) + token_view: Final = LiteLLM_VerificationTokenView(**response, last_refreshed_at=time.time()) # for prisma we need to cast the expires time to str - if response.expires is not None and isinstance(response.expires, datetime): - response.expires = response.expires.isoformat() + if token_view.expires is not None and isinstance(token_view.expires, datetime): + token_view.expires = token_view.expires.isoformat() + response = token_view return response except Exception as e: import traceback @@ -5989,13 +6081,12 @@ class PrismaClient: elif table_name == "user": db_data = self.jsonify_object(data=data) try: - new_user_row: Final = await UserRepository(self).table.upsert( - where={"user_id": data["user_id"]}, - data={ - "create": {**db_data}, - "update": {}, # don't do anything if it already exists - }, - ) + new_user_row: Final = await UserRepository(self).upsert_user(db_data) + if new_user_row is None: + raise HTTPException( + status_code=403, + detail="SSO for more than 5 users requires a LiteLLM Enterprise license. Existing users can still sign in.", + ) except Exception as e: if ( "Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`" @@ -6197,7 +6288,7 @@ class PrismaClient: """ Batch write update queries """ - batcher: _UpdateManyBatch = self.db.batch_() + batcher: _UpdateManyBatch = _update_many_batch(self.db) for idx, t in enumerate(data_list): # check if plain text or hash if t.token.startswith("sk-"): @@ -6222,7 +6313,7 @@ class PrismaClient: """ Batch write update queries """ - batcher = self.db.batch_() + batcher = _update_many_batch(self.db) for user in sorted(data_list, key=attrgetter("user_id")): try: data_json = self.jsonify_object(data=user.model_dump(exclude_none=True)) @@ -6247,7 +6338,7 @@ class PrismaClient: """ Batch write update queries """ - batcher = self.db.batch_() + batcher = _update_many_batch(self.db) for enduser in data_list: try: data_json = self.jsonify_object(data=enduser.model_dump(exclude_none=True)) @@ -6272,7 +6363,7 @@ class PrismaClient: """ Batch write update queries """ - batcher = self.db.batch_() + batcher = _update_many_batch(self.db) for budget in data_list: try: data_json = self.jsonify_object(data=budget.model_dump(exclude_none=True)) @@ -6295,7 +6386,7 @@ class PrismaClient: and isinstance(data_list, list) ): # Batch write update queries - batcher = self.db.batch_() + batcher = _update_many_batch(self.db) for idx, team in enumerate(data_list): try: data_json = self.jsonify_team_object(db_data=team.model_dump(exclude_none=True)) @@ -6960,7 +7051,7 @@ class PrismaClient: # caller's own retried query, a stronger check than SELECT 1, # and a reader that fails to come back sets `_reader_unavailable` # so reads fall through to the writer just recreated here. - await self.writer_db.query_raw("SELECT 1") + await _query_raw_rows(self.writer_db, "SELECT 1") # A recreate can decline: the optimistic-lock guard no-ops when # the writer generation moved since cycle entry, and the routing # wrapper then leaves the reader untouched as well. Callers that @@ -7308,7 +7399,7 @@ class PrismaClient: async def _view_setup_gate_table_present(self) -> bool: rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( - await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) + await _query_raw_rows(self.db, "SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) ) return rows[0]["present"] @@ -7318,7 +7409,7 @@ class PrismaClient: await asyncio.sleep(self._db_health_watchdog_interval_seconds) async with db_span("db_health_watchdog", None): await asyncio.wait_for( - self.db.query_raw("SELECT 1"), + _query_raw_rows(self.db, "SELECT 1"), timeout=self._db_health_watchdog_probe_timeout_seconds, ) if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable: @@ -7378,7 +7469,7 @@ class PrismaClient: async def _writer_is_read_only(self, writer: PrismaWrapper) -> bool: """True iff the pooled writer session answers reads but rejects writes (SQLSTATE 25006).""" rows: Final = _WRITER_WRITABILITY_PROBE_ROWS.validate_python( - await writer.query_raw(_WRITER_WRITABILITY_PROBE_SQL) + await _query_raw_rows(writer, _WRITER_WRITABILITY_PROBE_SQL) ) return any(row.get("transaction_read_only") == "on" for row in rows) @@ -7405,7 +7496,7 @@ class PrismaClient: """ sql_query: Final = "SELECT 1" async with db_span("health_check", None): - response: Final[object] = await wrapper.query_raw(sql_query) + response: Final[object] = await _query_raw_rows(wrapper, sql_query) return response async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool: @@ -7958,7 +8049,7 @@ class ProxyUpdateSpend: try: async with ( db_span("update_end_user_spend", "LiteLLM_EndUserTable"), - prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction, + _end_user_spend_transaction(prisma_client.db, timedelta(seconds=60)) as transaction, ): batcher: _EndUserSpendBatch async with transaction.batch_() as batcher: @@ -8101,6 +8192,18 @@ class ProxyUpdateSpend: return False +async def commit_tool_discoveries(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None: + """Scheduler job writing discovered tools to LiteLLM_ToolTable, on its own lock, next to update_spend.""" + writer: Final = proxy_logging_obj.db_spend_update_writer + await tool_registry_writer.commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=writer.tool_discovery_queue, + redis_update_buffer=writer.redis_update_buffer, + pod_lock_manager=writer.pod_lock_manager, + use_redis_buffer=RedisUpdateBuffer.should_commit_spend_updates_to_redis(), + ) + + async def update_spend( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index d4643994043..6d9249d5572 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -4,6 +4,7 @@ from typing import Final from fastapi import APIRouter, Depends, File, Form, Request, Response, UploadFile from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from litellm.proxy._types import * @@ -36,6 +37,7 @@ from litellm.types.videos.utils import ( ) router: Final = APIRouter() +_OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) @router.post( @@ -582,7 +584,9 @@ async def video_create_character( if video_file: data["video"] = video_file[0] - target_model_name: Final = extract_model_from_target_model_names(data.get("target_model_names")) + target_model_name: Final = extract_model_from_target_model_names( + _OPAQUE_VALUE_ADAPTER.validate_python(data.get("target_model_names")) + ) if target_model_name and not data.get("model"): data["model"] = target_model_name diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index a2be42f1b76..56123e4bee3 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: from litellm.router import Router -def extract_model_from_target_model_names(target_model_names: Any) -> str | None: +def extract_model_from_target_model_names(target_model_names: object) -> str | None: if isinstance(target_model_names, str): model_names: Final = [m.strip() for m in target_model_names.split(",") if m.strip()] return model_names[0] if model_names else None diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index 1e7aa6ac1c6..5520912dd41 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -106,11 +106,11 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): ) self.embedding_model = self.vector_store_config.get("embedding_model") or "amazon.titan-embed-text-v2:0" - self.wait_for_ingestion = self.vector_store_config.get("wait_for_ingestion", False) + self.wait_for_ingestion: object = self.vector_store_config.get("wait_for_ingestion", False) self.ingestion_timeout: int = _get_int(self.vector_store_config.get("ingestion_timeout"), 300) # Get AWS region using BaseAWSLLM method - _aws_region: Final = self.vector_store_config.get("aws_region_name") + _aws_region: Final[object] = self.vector_store_config.get("aws_region_name") self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( aws_region_name=str(_aws_region) if _aws_region else None ) diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index 4ee16685e03..6320e477bf2 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Final, cast from pydantic import ConfigDict from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, header_value, @@ -107,7 +108,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): raise ValueError("GEMINI_API_BASE is required") api_version: Final = "v1beta" - base_url: Final = f"{api_base}/{api_version}" + base_url: Final = join_api_base(api_base, api_version) # Create File Search store if not provided if not vector_store_id: @@ -148,7 +149,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): Returns: Store name (format: fileSearchStores/xxxxxxx) """ - url: Final = f"{base_url}/fileSearchStores" + url: Final = join_api_base(base_url, "fileSearchStores") request_body: Final = {"displayName": display_name} @@ -236,7 +237,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): # base_url is like: https://generativelanguage.googleapis.com/v1beta # We need: https://generativelanguage.googleapis.com/upload/v1beta/{store_id}:uploadToFileSearchStore api_base: Final = base_url.replace("/v1beta", "") # Get base without version - url: Final = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore" + url: Final = join_api_base(api_base, f"upload/v1beta/{vector_store_id}:uploadToFileSearchStore") # Build request body with chunking config and metadata if provided request_body: Final[dict[str, object]] = {"displayName": filename} diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 088dfb54155..7b9068ff95e 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -33,7 +33,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager, resolve_api_key_secret_reference -from ..litellm_core_utils.core_helpers import normalize_drop_params +from ..litellm_core_utils.core_helpers import as_str_mapping, 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 @@ -83,6 +83,13 @@ def _with_resolved_session_model(session: dict[str, object], model_name: str) -> return {**session, "model": model_name} +def _resolved_request_session(request_data: Mapping[str, object], model_name: str) -> Mapping[str, object] | None: + session: Final = as_str_mapping(request_data.get("session")) + if not isinstance(session, dict): + return None + return _with_resolved_session_model(session, model_name) + + def _build_litellm_metadata(kwargs: dict) -> dict[str, object]: """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" metadata: Final[dict[str, object]] = {**(kwargs.get("litellm_metadata") or {})} @@ -189,8 +196,9 @@ async def acreate_realtime_client_secret( custom_llm_provider=custom_llm_provider, ) request_data: Final = req.model_dump(exclude_none=True, exclude={"model"}) - if isinstance(request_data.get("session"), dict): - request_data["session"] = _with_resolved_session_model(request_data["session"], model_name) + resolved_session: Final = _resolved_request_session(request_data, model_name) + if resolved_session is not None: + request_data["session"] = resolved_session return await base_llm_http_handler.async_realtime_client_secret_handler( api_base=resolved_api_base, api_key=await _realtime_http_bearer_token( diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py index 81fba770b70..92f8ba984f4 100644 --- a/litellm/repositories/base_repository.py +++ b/litellm/repositories/base_repository.py @@ -8,6 +8,7 @@ from typing import Any, Final, Generic, Protocol, TypeVar, runtime_checkable from pydantic import BaseModel +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.repositories.chunked_in import find_many_in from litellm.repositories.prisma_protocols import TableActions @@ -131,4 +132,4 @@ def is_unique_violation(exc: BaseException) -> bool: return "P2002" in str(exc) or "unique constraint" in str(exc).lower() if isinstance(exc, UniqueViolationError): return True - return getattr(exc, "code", None) == "P2002" + return attribute_or_none(exc, "code") == "P2002" diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index c42301a9316..63cfcbc7b45 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -171,4 +171,6 @@ class PrismaBatch(Protocol): @property def litellm_projecttable(self) -> BatchTable: ... + def execute_raw(self, query: str, *args: object) -> None: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 57f6fd33c11..b89d87aa354 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -75,6 +75,7 @@ class _PrismaClientView(Protocol): _MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member]) +_RAW_COLUMN: Final = TypeAdapter(object) _JSON_ENCODED_TEAM_FIELDS: Final = ( "metadata", "model_spend", @@ -142,7 +143,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): ) if not rows: return None - raw_value: Final[object] = rows[0]["members_with_roles"] + raw_value: Final = _RAW_COLUMN.validate_python(rows[0]["members_with_roles"]) parsed: Final[object] = json.loads(raw_value) if isinstance(raw_value, str) else raw_value if not parsed: return [] diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 95f4f5e81be..f6aad883479 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -15,11 +15,11 @@ instead of full-model writes, which trip ``prisma.errors.DataError`` on rows carrying fields the update input type rejects (see #27730). """ -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime -from typing import Final +from typing import Final, Protocol, runtime_checkable from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch @@ -62,6 +62,13 @@ class TeamSpendResetWrites: ) +@runtime_checkable +class SpendResetWrites(Protocol): + def queue_spend_zero(self, where: Mapping[str, object]) -> None: ... + + def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None: ... + + @dataclass(frozen=True, slots=True) class LinkedSpendResetWrites: table: BatchTable @@ -78,6 +85,30 @@ class LinkedSpendResetWrites: ) +# Membership rows written before total_spend existed (migration 20260421135425, default 0) hold their +# history only in spend. Zeroing spend would erase it, so the floor lands in the same batch, ahead of the reset. +_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL: Final = """ +UPDATE "LiteLLM_TeamMembership" +SET total_spend = spend +WHERE budget_id = ANY($1::text[]) AND spend > total_spend +""" + + +@dataclass(frozen=True, slots=True) +class TeamMembershipSpendResetWrites: + resets: LinkedSpendResetWrites + batch: PrismaBatch + + def queue_spend_zero(self, where: Mapping[str, object]) -> None: + self.resets.queue_spend_zero(where) + + def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None: + self.resets.queue_spend_decrement(where, amount) + + def queue_total_spend_repair(self, budget_ids: Sequence[str]) -> None: + self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids)) + + @dataclass(frozen=True, slots=True) class BudgetWindowWrites: table: BatchTable @@ -106,7 +137,7 @@ class BudgetCascadeUnitOfWork: ``enduser_spend_reset_unit_of_work`` because one tier can gate millions. """ - team_memberships: LinkedSpendResetWrites + team_memberships: TeamMembershipSpendResetWrites keys: LinkedSpendResetWrites organizations: LinkedSpendResetWrites tags: LinkedSpendResetWrites @@ -132,7 +163,9 @@ async def budget_cascade_unit_of_work( ) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]: batch = new_batch() yield BudgetCascadeUnitOfWork( - team_memberships=LinkedSpendResetWrites(table=batch.litellm_teammembership), + team_memberships=TeamMembershipSpendResetWrites( + resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch + ), keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken), organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 5ed1574c6bc..1fce0717ea3 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -2,8 +2,11 @@ User repository for database operations on LiteLLM_UserTable. """ +import asyncio import json -from collections.abc import Mapping, Sequence +from collections.abc import Generator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, contextmanager +from contextvars import ContextVar from itertools import chain from typing import TYPE_CHECKING, Final, Protocol @@ -35,6 +38,16 @@ ORDER BY p.user_id """ _PLACEHOLDER_ROWS_ADAPTER: Final = TypeAdapter(tuple[SCIMPlaceholder, ...]) +_FREE_SSO_CREATION: Final[ContextVar[tuple[str | None, object] | None]] = ContextVar("free_sso_creation", default=None) + + +@contextmanager +def free_sso_user_creation(user_id: str | None, *, enabled: bool = True) -> Generator[None]: + token: Final = _FREE_SSO_CREATION.set((user_id, asyncio.current_task()) if enabled else None) + try: + yield + finally: + _FREE_SSO_CREATION.reset(token) class _UserDb(Protocol): @@ -42,6 +55,12 @@ class _UserDb(Protocol): def litellm_usertable(self) -> TableActions["prisma_models.LiteLLM_UserTable"]: ... +class _UserTransaction(_UserDb, Protocol): + async def execute_raw(self, query: str) -> int: ... + + async def query_raw(self, query: str) -> object: ... + + class _PrismaClientView(Protocol): @property def db(self) -> _UserDb: ... @@ -49,6 +68,16 @@ class _PrismaClientView(Protocol): @property def writer_db(self) -> _UserDb: ... + def tx(self) -> AbstractAsyncContextManager[_UserTransaction]: ... + + +async def _count_billable_users(table: TableActions["prisma_models.LiteLLM_UserTable"]) -> int: + from prisma import Json + + total: Final = await table.count() + deactivated: Final = await table.count(where={"metadata": {"path": ["scim_active"], "equals": Json(False)}}) + return max(0, total - deactivated) + class UserRepository(BaseRepository[LiteLLM_UserTable]): """Repository for user database operations.""" @@ -137,11 +166,32 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): null, or true all count, so seats are counted as total users minus the deactivated ones. """ - from prisma import Json # pyright: ignore[reportUnknownVariableType] + return await _count_billable_users(self.table) - total: Final = await self.count() - deactivated: Final = await self.count(where={"metadata": {"path": ["scim_active"], "equals": Json(False)}}) - return max(0, total - deactivated) + async def upsert_user(self, data: Mapping[str, object]) -> "prisma_models.LiteLLM_UserTable | None": + """Return the user row, or None when a new free SSO user would exceed capacity.""" + scope: Final = _FREE_SSO_CREATION.get() + if ( + scope is None + or (scope[0] is not None and scope[0] != data["user_id"]) + or scope[1] is not asyncio.current_task() + ): + return await self.table.upsert( + where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}} + ) + + client: Final[_PrismaClientView] = self.prisma_client + async with client.tx() as transaction: + await transaction.execute_raw("SET TRANSACTION ISOLATION LEVEL READ COMMITTED") + await transaction.query_raw("SELECT 1 AS locked FROM pg_advisory_xact_lock(468530001)") + existing: Final = await transaction.litellm_usertable.find_unique(where={"user_id": data["user_id"]}) + if existing is not None: + return existing + if await _count_billable_users(transaction.litellm_usertable) >= 5: + return None + return await transaction.litellm_usertable.upsert( + where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}} + ) async def create_user( self, diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index f7e2616324f..65d9b89cf79 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -1,6 +1,6 @@ import asyncio import contextvars -from collections.abc import Awaitable +from collections.abc import Awaitable, Coroutine from functools import partial from typing import Final, Literal @@ -76,7 +76,7 @@ async def arerank( func_with_context: Final = partial(ctx.run, func) init_response: Final = await loop.run_in_executor(None, func_with_context) - if asyncio.iscoroutine(init_response): + if isinstance(init_response, Coroutine): response = await init_response else: response = init_response diff --git a/litellm/rerank_api/rerank_utils.py b/litellm/rerank_api/rerank_utils.py index f365ae749d1..4924625b011 100644 --- a/litellm/rerank_api/rerank_utils.py +++ b/litellm/rerank_api/rerank_utils.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -8,7 +8,7 @@ def get_optional_rerank_params( model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 60318d6b184..8b00cdf4813 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -175,7 +175,7 @@ def responses_api_bridge_check( pinned_mode: Final = ( None if explicit_responses_prefix else deployment_routing_mode(deployment_model_info, custom_llm_provider) ) - if explicit_responses_prefix or pinned_mode == "responses": + if custom_llm_provider == LlmProviders.CHATGPT or explicit_responses_prefix or pinned_mode == "responses": model = model.replace("responses/", "") model_info["mode"] = "responses" elif pinned_mode == "chat" and backend_serves_chat_completions(model_info): diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 7121908517b..f29f6c83ef0 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -15,7 +15,7 @@ import json import time import uuid from collections.abc import Iterable, Sequence -from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast # noqa: TID251 # see kwargs-ok / cast-ok markers +from typing import TYPE_CHECKING, Final, TypeAlias, cast # noqa: TID251 # see cast-ok markers from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -517,7 +517,7 @@ async def aresponses_with_emulated_file_search( model: str, tools: Iterable[ToolParam] | None = None, # Pass-through params — forwarded as-is to the underlying aresponses call - **kwargs: Any, # kwargs-ok: `object` would surface the caller's partially-unknown dict at its call site + **kwargs: object, # kwargs-ok: pass-through params forwarded as-is to the underlying aresponses call ) -> ResponsesAPIResponse: """ Emulated file_search for providers that don't support it natively. diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py new file mode 100644 index 00000000000..8bd82c49090 --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -0,0 +1,329 @@ +import json +import uuid +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain, groupby +from types import MappingProxyType +from typing import Final, TypeVar + +from openai.types.responses.easy_input_message_param import EasyInputMessageParam +from openai.types.responses.response_compaction_item import ResponseCompactionItem +from openai.types.responses.response_create_params import ResponseInputParam +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm.responses.utils import ResponsesAPIRequestUtils + +from .tool_search import chat_config_builds_anthropic_request, provider_chat_config + +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST: Final = TypeAdapter(list[object]) +_Item: Final = TypeVar("_Item") +_INSTRUCTION_ROLES: Final = frozenset(("system", "developer")) + + +def _json_object(value: object) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + + +@dataclass(frozen=True, slots=True) +class CompactionBlock: + summary: str + wire_block: Mapping[str, object] + + +def _compaction_block(value: object) -> CompactionBlock | None: + block: Final = _json_object(value) + if block is None: + return None + summary: Final = block.get("content") + if block.get("type") != "compaction" or not isinstance(summary, str) or not summary: + return None + return CompactionBlock(summary=summary, wire_block=block) + + +def _newest_compaction_block(compaction_blocks: object) -> CompactionBlock | None: + if not isinstance(compaction_blocks, list): + return None + blocks: Final = _OBJECT_LIST.validate_python(compaction_blocks) + return next((block for value in reversed(blocks) if (block := _compaction_block(value)) is not None), None) + + +def _provider_compaction_block(provider_specific_fields: object) -> CompactionBlock | None: + fields: Final = _json_object(provider_specific_fields) + return None if fields is None else _newest_compaction_block(fields.get("compaction_blocks")) + + +def provider_compaction_summary(provider_specific_fields: object) -> str | None: + block: Final = _provider_compaction_block(provider_specific_fields) + return None if block is None else block.summary + + +class _AssistantTurnFields(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + content: object = None + tool_calls: object = None + provider_specific_fields: object = None + + +def carries_only_compaction(message: object) -> bool: + fields: Final = _AssistantTurnFields.model_validate(message) + return ( + not fields.content + and not fields.tool_calls + and provider_compaction_summary(fields.provider_specific_fields) is not None + ) + + +def _choice_provider_specific_fields(choice: object) -> object: + choice_fields: Final = _json_object(choice) + message: Final = None if choice_fields is None else _json_object(choice_fields.get("message")) + return None if message is None else message.get("provider_specific_fields") + + +def response_holds_compaction_block(response: object) -> bool: + fields: Final = _json_object(response) + if fields is None: + return False + if fields.get("object") == "response": + return holds_compaction_summary(fields.get("output")) + try: + choices: Final = _OBJECT_LIST.validate_python(fields.get("choices")) + except ValidationError: + return False + return any(provider_compaction_summary(_choice_provider_specific_fields(choice)) is not None for choice in choices) + + +def compaction_output_item(provider_specific_fields: object) -> ResponseCompactionItem | None: + block: Final = _provider_compaction_block(provider_specific_fields) + if block is None: + return None + return ResponseCompactionItem( + id=f"cmp_{uuid.uuid4()}", + encrypted_content=json.dumps(dict(block.wire_block), separators=(",", ":")), + type="compaction", + ) + + +def _entry_type(entry: object) -> object: + fields: Final = _json_object(entry) + return None if fields is None else fields.get("type") + + +def _has_entry_of_type(entries: object, entry_type: str) -> bool: + try: + rows: Final = _OBJECT_LIST.validate_python(entries) + except ValidationError: + return False + return any(_entry_type(row) == entry_type for row in rows) + + +def requests_compaction(context_management: object) -> bool: + if _has_entry_of_type(context_management, "compaction"): + return True + anthropic_form: Final = _json_object(context_management) + return anthropic_form is not None and _has_entry_of_type(anthropic_form.get("edits"), "compact_20260112") + + +def provider_takes_compaction_blocks(model: str, custom_llm_provider: str | None) -> bool: + from litellm.llms.bedrock.claude_platform.transformation import BedrockClaudePlatformConfig + + chat_config: Final = provider_chat_config(model, custom_llm_provider) + return chat_config_builds_anthropic_request(chat_config) and not isinstance( + chat_config, BedrockClaudePlatformConfig + ) + + +def _without_affinity_wrap(encrypted_content: str) -> str: + _, original = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(encrypted_content) + return original + + +def decode_compaction_item(encrypted_content: object) -> CompactionBlock | None: + if not isinstance(encrypted_content, str): + return None + try: + decoded: Final = _JSON_OBJECT.validate_json(_without_affinity_wrap(encrypted_content)) + except ValidationError: + return None + return _compaction_block(decoded) + + +def compaction_summary(item: object) -> str | None: + fields: Final = _json_object(item) + if fields is None or fields.get("type") != "compaction": + return None + block: Final = decode_compaction_item(fields.get("encrypted_content")) + return None if block is None else block.summary + + +def summary_only_compaction_content(summary: str) -> str: + return json.dumps({"type": "compaction", "content": summary}, separators=(",", ":")) + + +def with_compaction_summary(encrypted_content: str, summary: str) -> str: + original: Final = _without_affinity_wrap(encrypted_content) + block: Final = decode_compaction_item(original) + if block is None: + return encrypted_content + rewritten: Final = json.dumps({**block.wire_block, "content": summary}, separators=(",", ":")) + return encrypted_content.removesuffix(original) + rewritten + + +def with_visited_compaction_summary(item: object, visit: Callable[[str], str]) -> Mapping[str, object] | None: + fields: Final = _json_object(item) + if fields is None: + return None + summary: Final = compaction_summary(fields) + encrypted_content: Final = fields.get("encrypted_content") + if summary is None or not isinstance(encrypted_content, str): + return None + return {**fields, "encrypted_content": with_compaction_summary(encrypted_content, visit(summary))} + + +def holds_compaction_summary(input: object) -> bool: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return False + return any(compaction_summary(item) is not None for item in items) + + +def _is_instruction_item(item: object) -> bool: + fields: Final = _json_object(item) + if fields is None or fields.get("type", "message") != "message": + return False + role: Final = fields.get("role") + return isinstance(role, str) and role in _INSTRUCTION_ROLES + + +def input_items(input: object) -> tuple[object, ...]: + try: + return tuple(_OBJECT_LIST.validate_python(input)) + except ValidationError: + return () + + +def instruction_items(input: object) -> tuple[object, ...]: + return tuple(item for item in input_items(input) if _is_instruction_item(item)) + + +def _call_id(item: object) -> str | None: + fields: Final = _json_object(item) + call_id: Final = None if fields is None else fields.get("call_id") + return call_id if isinstance(call_id, str) else None + + +def _is_turn_item(item: object) -> bool: + fields: Final = _json_object(item) + item_type: Final = None if fields is None else fields.get("type", "message") + if fields is None or not isinstance(item_type, str): + return False + if item_type == "message": + return fields.get("role") == "assistant" + return item_type == "reasoning" or (_call_id(fields) is not None and not item_type.endswith("_output")) + + +def _turn_indices_answered_later( + turn: tuple[tuple[int, object], ...], calls_answered_later: frozenset[int] +) -> tuple[int, ...]: + if not any(index in calls_answered_later for index, _ in turn): + return () + return tuple(index for index, item in turn if _call_id(item) is None or index in calls_answered_later) + + +def _is_turn_entry(entry: tuple[int, object]) -> bool: + return _is_turn_item(entry[1]) + + +def _calls_answered_only_later(earlier: Sequence[object], answered: frozenset[str]) -> frozenset[int]: + last_index_by_call_id: Final = MappingProxyType( + {call_id: index for index, item in enumerate(earlier) if (call_id := _call_id(item)) is not None} + ) + return frozenset(index for call_id, index in last_index_by_call_id.items() if call_id in answered) + + +def _earlier_turn_indices_answered_later(earlier: Sequence[object], answered: frozenset[str]) -> tuple[int, ...]: + calls_answered_later: Final = _calls_answered_only_later(earlier, answered) + turns: Final = (tuple(turn) for is_turn, turn in groupby(enumerate(earlier), key=_is_turn_entry) if is_turn) + return tuple(chain.from_iterable(_turn_indices_answered_later(turn, calls_answered_later) for turn in turns)) + + +def _call_ids(items: Sequence[object]) -> frozenset[str]: + return frozenset(call_id for item in items if (call_id := _call_id(item)) is not None) + + +def _keeps_every_item(item: object) -> bool: + return True + + +def _with_answered_turns_after_newest_compaction( + input: Sequence[_Item], keeps_earlier_item: Callable[[object], bool] +) -> tuple[_Item, ...]: + newest: Final = next( + (index for index in range(len(input) - 1, 0, -1) if compaction_summary(input[index]) is not None), 0 + ) + later: Final = input[newest:] + moved: Final = _earlier_turn_indices_answered_later(input[:newest], _call_ids(later)) + moved_indices: Final = frozenset(moved) + moved_call_ids: Final = _call_ids([input[index] for index in moved]) + first_answer: Final = next( + (index for index, item in enumerate(later) if _call_id(item) in moved_call_ids), len(later) + ) + return ( + *(item for index, item in enumerate(input[:newest]) if index not in moved_indices and keeps_earlier_item(item)), + *later[:first_answer], + *(input[index] for index in moved), + *later[first_answer:], + ) + + +def _items_from_newest_compaction(input: Sequence[_Item]) -> tuple[_Item, ...]: + return _with_answered_turns_after_newest_compaction(input, _is_instruction_item) + + +def with_answered_turns_after_newest_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: + return ( + list(_with_answered_turns_after_newest_compaction(input, _keeps_every_item)) + if isinstance(input, list) + else input + ) + + +def without_items_before_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: + return list(_items_from_newest_compaction(input)) if isinstance(input, list) else input + + +def _summary_or_item(item: object) -> object: + summary: Final = compaction_summary(item) + return item if summary is None else summary + + +def billable_input_since_newest_compaction(input: object) -> object: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return input + return [_summary_or_item(item) for item in _items_from_newest_compaction(items)] + + +def _as_summary_message(item: _Item) -> _Item | EasyInputMessageParam: + summary: Final = compaction_summary(item) + return item if summary is None else EasyInputMessageParam(role="user", content=summary, type="message") + + +def with_compaction_items_as_summaries(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list) or not holds_compaction_summary(input): + return input + return [_as_summary_message(item) for item in _items_from_newest_compaction(input)] + + +def messages_since_newest_compaction(input: object) -> object: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return input + return [_as_summary_message(item) for item in _items_from_newest_compaction(items)] diff --git a/litellm/responses/litellm_completion_transformation/count_request.py b/litellm/responses/litellm_completion_transformation/count_request.py index 56a0b9024c5..7679b4ede71 100644 --- a/litellm/responses/litellm_completion_transformation/count_request.py +++ b/litellm/responses/litellm_completion_transformation/count_request.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from typing import Final, cast # noqa: TID251 # input items are forwarded as sent, as inference does @@ -14,7 +14,11 @@ from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body from litellm.responses.additional_tools import hoist_additional_tools from litellm.responses.file_search.emulated_handler import replace_file_search_tools -from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig +from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary +from litellm.responses.litellm_completion_transformation.transformation import ( + ChatCompletionSession, + LiteLLMCompletionResponsesConfig, +) from litellm.responses.main import previous_response_client_ids from litellm.types.llms.openai import ResponseInputParam from litellm.utils import ProviderConfigManager @@ -132,12 +136,17 @@ def _with_emulated_file_search(responses_request: Mapping[str, object]) -> Mappi return {**responses_request, "tools": replace_file_search_tools(_ITEMS.validate_python(tools))[0]} +def _with_compaction_as_summary_text(responses_request: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in responses_request.items() if name != "context_management"} + + async def _with_previous_turns( chat_body: Mapping[str, object], model: str, custom_llm_provider: str, responses_request: Mapping[str, object], policy: DropPolicy, + session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None, ) -> Mapping[str, object]: previous_response_id: Final = responses_request.get("previous_response_id") if not isinstance(previous_response_id, str) or not previous_response_id: @@ -147,8 +156,10 @@ async def _with_previous_turns( previous_response_id=previous_response_id, litellm_completion_request={**chat_body, "model": model, "custom_llm_provider": custom_llm_provider}, instructions=instructions if isinstance(instructions, str) else None, + session_lookup=session_lookup, request_kwargs={"drop_params": policy.drop_params}, previous_response_client_ids=previous_response_client_ids(previous_response_id, responses_request), + input_replaces_history=holds_compaction_summary(responses_request.get("input")), ) @@ -171,6 +182,7 @@ async def chat_count_request( api_base: str | None, responses_request: Mapping[str, object], policy: DropPolicy, + session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, ) -> ChatCountRequest | ChatCountBridgeError | None: """The chat request Responses inference sends a provider with no native Responses API, or None when inference sends this provider the Responses body itself.""" @@ -184,7 +196,9 @@ async def chat_count_request( try: chat_body: Final = await _with_previous_turns( responses_batch_body_to_chat_body( - _with_emulated_file_search(_with_hoisted_additional_tools(responses_request)), + _with_compaction_as_summary_text( + _with_emulated_file_search(_with_hoisted_additional_tools(responses_request)) + ), custom_llm_provider, policy.drop_params, ), @@ -192,6 +206,7 @@ async def chat_count_request( custom_llm_provider, responses_request, policy, + session_lookup, ) tools: Final = ( () if "tools" in policy.additional_drop_params else tuple(map(_unmarked, _mappings(chat_body.get("tools")))) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index fcda6b31b00..98e1984550c 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -23,6 +23,7 @@ from typing import Final, Literal, TypeAlias from pydantic import TypeAdapter, ValidationError from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( ChatCompletionToolParam, @@ -30,6 +31,7 @@ from litellm.types.llms.openai import ( ) _MAX_ARGUMENTS_LEN: Final = 1_000_000 +_DECODED_JSON: Final = TypeAdapter(object) TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE: Final = MappingProxyType( {"function_call": "fc", "custom_tool_call": "ctc", "tool_search_call": "tsc"} @@ -111,8 +113,8 @@ def unwrap_custom_tool_arguments(arguments: str) -> str: if len(arguments) > _MAX_ARGUMENTS_LEN: return arguments try: - parsed: Final = json.loads(arguments) - if isinstance(parsed, dict) and "content" in parsed: + parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(arguments))) + if parsed is not None and "content" in parsed: return str(parsed["content"]) except (json.JSONDecodeError, TypeError, ValueError): pass diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 5489d2223ac..c9d800421f7 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -9,11 +9,13 @@ from pydantic import TypeAdapter import litellm from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, + with_compaction_items_as_params, ) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( @@ -120,6 +122,7 @@ class LiteLLMCompletionTransformationHandler: instructions=responses_api_request.get("instructions"), request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), previous_response_client_ids=previous_response_client_ids, + input_replaces_history=holds_compaction_summary(with_compaction_items_as_params(request_input)), ) acompletion_args: Final = {} diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 35160830fdc..d50e9a61edc 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,12 +1,20 @@ import asyncio import json from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import chain from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler +from litellm.responses.litellm_completion_transformation.compaction_items import ( + holds_compaction_summary, + input_items, + instruction_items, + response_holds_compaction_block, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( AllMessageValues, @@ -20,6 +28,7 @@ if TYPE_CHECKING: from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.responses.litellm_completion_transformation.transformation import ( ChatCompletionSession, + SessionMessage, ) else: ChatCompletionSession = Any @@ -66,6 +75,95 @@ def _stored_responses_api_output(stored_response: object) -> ResponseInputParam return None +@dataclass(frozen=True, slots=True) +class _ReplayStart: + index: int + at_output: bool + + +def _compaction_point( + index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None +) -> _ReplayStart | None: + if response_holds_compaction_block(spend_log.get("response")): + return _ReplayStart(index=index, at_output=True) + if proxy_server_request is not None and holds_compaction_summary(proxy_server_request.get("input")): + return _ReplayStart(index=index, at_output=False) + return None + + +def _replay_start( + spend_logs: Sequence["SpendLogsPayload"], proxy_server_requests: Sequence[Mapping[str, object] | None] +) -> _ReplayStart: + return next( + ( + start + for index in range(len(spend_logs) - 1, -1, -1) + if (start := _compaction_point(index, spend_logs[index], proxy_server_requests[index])) is not None + ), + _ReplayStart(index=0, at_output=False), + ) + + +def _holds_compaction_point( + spend_logs: Sequence["SpendLogsPayload"], proxy_server_requests: Sequence[Mapping[str, object] | None] +) -> bool: + return any( + _compaction_point(index, spend_log, proxy_server_request) is not None + for index, (spend_log, proxy_server_request) in enumerate(zip(spend_logs, proxy_server_requests, strict=True)) + ) + + +def _referenced_request_ids(previous_response_id: str, client_response_ids: tuple[str, ...]) -> tuple[str, ...]: + decoded_response_id: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(previous_response_id) + response_id: Final = decoded_response_id.get("response_id", previous_response_id) + return tuple(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) + + +def _turns_through_referenced_response( + spend_logs: Sequence["SpendLogsPayload"], referenced_request_ids: tuple[str, ...] +) -> int: + return next( + ( + index + 1 + for index in range(len(spend_logs) - 1, -1, -1) + if spend_logs[index].get("request_id") in referenced_request_ids + ), + len(spend_logs), + ) + + +def _replayed_request( + index: int, + proxy_server_request: Mapping[str, object] | None, + start: _ReplayStart, + earlier_instructions: tuple[object, ...], +) -> Mapping[str, object] | None: + if proxy_server_request is None or index != start.index: + return proxy_server_request + stored_input: Final = proxy_server_request.get("input") + if start.at_output: + return { + **proxy_server_request, + "input": [*earlier_instructions, *instruction_items(stored_input)], + "messages": None, + } + if not earlier_instructions: + return proxy_server_request + return {**proxy_server_request, "input": [*earlier_instructions, *input_items(stored_input)]} + + +def _instruction_items_before( + proxy_server_requests: Sequence[Mapping[str, object] | None], start: _ReplayStart +) -> tuple[object, ...]: + return tuple( + chain.from_iterable( + instruction_items(request.get("input")) + for request in proxy_server_requests[: start.index] + if request is not None + ) + ) + + def _stored_instructions(proxy_server_request: Mapping[str, object] | None) -> str | None: instructions: Final = None if proxy_server_request is None else proxy_server_request.get("instructions") return instructions if isinstance(instructions, str) and instructions else None @@ -91,6 +189,7 @@ class ResponsesSessionHandler: """ from litellm.responses.litellm_completion_transformation.transformation import ( ChatCompletionSession, + with_stored_calls_the_input_answers, ) verbose_proxy_logger.debug("inside get_chat_completion_message_history_for_previous_response_id") @@ -107,26 +206,34 @@ class ResponsesSessionHandler: if len(all_spend_logs) > 0: litellm_session_id = all_spend_logs[0].get("session_id") - chat_completion_message_history: list[ - AllMessageValues - | GenericChatCompletionMessage - | ChatCompletionMessageToolCall - | ChatCompletionResponseMessage - | Message - ] = [] proxy_server_requests: Final = [ await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) for spend_log in all_spend_logs ] - for spend_log, proxy_server_request_dict in zip(all_spend_logs, proxy_server_requests): - chat_completion_message_history = ( - ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( - spend_log=spend_log, - proxy_server_request_dict=proxy_server_request_dict, - chat_completion_message_history=chat_completion_message_history, - item_cache_control=item_cache_control, - ) - ) + replayable_turns: Final = _turns_through_referenced_response( + all_spend_logs, _referenced_request_ids(previous_response_id, client_response_ids) + ) + start: Final = _replay_start(all_spend_logs[:replayable_turns], proxy_server_requests[:replayable_turns]) + replayed_turns: Final = ( + replayable_turns if _holds_compaction_point(all_spend_logs, proxy_server_requests) else len(all_spend_logs) + ) + earlier_instructions: Final = _instruction_items_before(proxy_server_requests, start) + replayed_history: Final = ResponsesSessionHandler._history_of_turns( + all_spend_logs[start.index : replayed_turns], + [ + _replayed_request(index, proxy_server_requests[index], start, earlier_instructions) + for index in range(start.index, replayed_turns) + ], + item_cache_control, + ) + history_before_replay: Final = ResponsesSessionHandler._history_of_turns( + all_spend_logs[: start.index], proxy_server_requests[: start.index], item_cache_control + ) + chat_completion_message_history: Final = ( + list(with_stored_calls_the_input_answers(replayed_history, history_before_replay)) + if history_before_replay + else replayed_history + ) verbose_proxy_logger.debug( "chat_completion_message_history %s", @@ -138,6 +245,24 @@ class ResponsesSessionHandler: instructions=_latest_replayable_instructions(proxy_server_requests), ) + @staticmethod + def _history_of_turns( + spend_logs: Sequence["SpendLogsPayload"], + proxy_server_requests: Sequence[Mapping[str, object] | None], + item_cache_control: bool, + ) -> list["SessionMessage"]: + return list( + chain.from_iterable( + ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( + spend_log=spend_log, + proxy_server_request_dict=proxy_server_request, + chat_completion_message_history=[], + item_cache_control=item_cache_control, + ) + for spend_log, proxy_server_request in zip(spend_logs, proxy_server_requests, strict=True) + ) + ) + @staticmethod def extend_chat_completion_message_with_spend_log_payload( spend_log: "SpendLogsPayload", @@ -182,6 +307,7 @@ class ResponsesSessionHandler: responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, + only_since_compaction=True, ) ) @@ -196,6 +322,7 @@ class ResponsesSessionHandler: responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, + only_since_compaction=True, ) ) elif isinstance(_response_output, dict) and _response_output and _response_output != {}: @@ -342,12 +469,11 @@ class ResponsesSessionHandler: verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) - decoded_response_id: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(previous_response_id) - response_id: Final = decoded_response_id.get("response_id", previous_response_id) + request_ids: Final = list(_referenced_request_ids(previous_response_id, client_response_ids)) + response_id: Final = request_ids[0] if prisma_client is None: return [] - request_ids: Final = list(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) query: Final = """ WITH matching_session AS ( SELECT session_id, api_key, "user", team_id diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 84d4648f75f..b5ef9e95741 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -5,8 +5,11 @@ from itertools import filterfalse from typing import Any, Final, cast import litellm -from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items +from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.main import stream_chunk_builder +from litellm.responses.litellm_completion_transformation.compaction_items import compaction_output_item from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, extract_custom_tool_names, @@ -82,9 +85,18 @@ def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | ) +def _is_signed_thinking_block(block: object) -> bool: + fields: Final = as_str_mapping(block) + return fields is not None and isinstance(block, dict) and bool(fields.get("signature") or fields.get("data")) + + +def _streams_compaction(chunk: ModelResponseStream) -> bool: + delta_fields: Final = chunk.choices[0].delta.provider_specific_fields if chunk.choices else None + return delta_fields is not None and "compaction_blocks" in delta_fields + + def _delta_has_signed_thinking_block(delta: object) -> bool: - blocks: Final = getattr(delta, "thinking_blocks", None) or () - return any(isinstance(b, dict) and (b.get("signature") or b.get("data")) for b in blocks) + return any(_is_signed_thinking_block(block) for block in attribute_items(delta, "thinking_blocks")) def _delta_carries_output(delta: ChatCompletionDelta) -> bool: @@ -145,6 +157,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id: str | None = None self._message_output_index: int | None = None self._reasoning_output_index: int | None = None + self._compaction_output_index: int | None = None + self._compaction_item_id: str | None = None + self._sent_compaction_item_done: bool = False self._cached_response_id: str | None = None self._buffered_chunk: ModelResponseStream | None = None self._upstream_exhausted: bool = False @@ -185,7 +200,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_search_call_ids = self._tool_search_call_ids | {call_id} return True - def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: + def _next_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> OutputItemAddedEvent: self._sequence_number += 1 added: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -193,7 +208,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): item=item, ) added.__dict__["sequence_number"] = self._sequence_number - self._pending_tool_events.append(added) + return added + + def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: + self._pending_tool_events.append(self._next_output_item_added(output_index, item)) def _queue_tool_search_call_added(self, call_id: str, output_index: int, args_delta: str) -> None: if call_id not in self._tool_args_by_call_id: @@ -247,6 +265,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return self._message_output_index case "reasoning": return self._reasoning_output_index + case "compaction": + return self._compaction_output_index case _: call_id: Final = getattr(item, "call_id", None) or str(getattr(item, "id", "")).removeprefix("ws_") return self._tool_output_index_by_call_id.get(str(call_id)) @@ -255,6 +275,45 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): index: Final = self._streamed_output_index(item) return (index is None, index or 0) + def _streamed_compaction_item(self) -> BaseLiteLLMOpenAIResponseObject | None: + item: Final = compaction_output_item(self._accumulated_provider_specific_fields) + if item is None: + return None + self._compaction_item_id = self._compaction_item_id or item.id + return BaseLiteLLMOpenAIResponseObject(**{**item.model_dump(exclude_none=True), "id": self._compaction_item_id}) + + def _queue_compaction_item_added(self) -> None: + if ( + self._compaction_output_index is not None + or "compaction_blocks" not in self._accumulated_provider_specific_fields + ): + return + streamed_item: Final = self._streamed_compaction_item() + if streamed_item is None: + return + self._compaction_output_index = self._allocate_output_index() + unfinished_item: Final = streamed_item.model_copy(update={"encrypted_content": ""}) + self._pending_response_events.append( + self._next_output_item_added(self._compaction_output_index, unfinished_item) + ) + + def _queue_compaction_item_done(self) -> None: + if self._compaction_output_index is None or self._sent_compaction_item_done: + return + streamed_item: Final = self._streamed_compaction_item() + if streamed_item is None: + return + self._sent_compaction_item_done = True + self._sequence_number += 1 + self._pending_response_events.append( + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=self._compaction_output_index, + sequence_number=self._sequence_number, + item=streamed_item, + ) + ) + def _normalize_tool_call_index(self, tool_call: object) -> int | None: idx_raw: Final = tool_call.get("index") if isinstance(tool_call, dict) else getattr(tool_call, "index", None) if idx_raw is None: @@ -591,6 +650,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ) + def _ordered_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + ordered: Final = self._reasoning_part_first(chunk) + if not _streams_compaction(ordered): + self._queue_compaction_item_added() + return ordered + def _with_encoded_response_id(self, response: ResponsesAPIResponse) -> ResponsesAPIResponse: return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, @@ -1003,6 +1068,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if not self.litellm_model_response or isinstance(self.litellm_model_response, TextCompletionResponse): self.litellm_model_response = self.create_litellm_model_response() if self.litellm_model_response: + self._queue_compaction_item_added() + self._queue_compaction_item_done() + if self._pending_response_events: + return self._pending_response_events.pop(0) # If tool calls exist, emit tool events before finishing/response.completed. if isinstance(self.litellm_model_response, ModelResponse): self._queue_final_tool_call_done_events(self.litellm_model_response) @@ -1090,6 +1159,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._reasoning_item_id or self._cached_reasoning_item_id or mint_reasoning_item_id() ) reasoning_item_id: Final = self._cached_reasoning_item_id + self._queue_compaction_item_done() self._sequence_number += 1 text_done_event: Final = self.create_reasoning_summary_text_done_event( reasoning_item_id=reasoning_item_id, @@ -1138,7 +1208,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): raise StopAsyncIteration chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: - chunk = self._reasoning_part_first(cast(ModelResponseStream, chunk)) + chunk = self._ordered_chunk(cast(ModelResponseStream, chunk)) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1194,11 +1264,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): try: buffered_chunk = self._take_buffered_chunk() if buffered_chunk is not None: - chunk = self._reasoning_part_first(buffered_chunk) + chunk = self._ordered_chunk(buffered_chunk) elif self._upstream_exhausted: raise StopIteration else: - chunk = self._reasoning_part_first(self.litellm_custom_stream_wrapper.__next__()) + chunk = self._ordered_chunk(self.litellm_custom_stream_wrapper.__next__()) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1355,13 +1425,25 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "reasoning", self._cached_reasoning_item_id, ) - streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, reasoning_aligned) + compaction_aligned: Final = _output_items_with_id(reasoning_aligned, "compaction", self._compaction_item_id) + streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, compaction_aligned) return tuple(sorted(streamed_items, key=self._streamed_output_position)) def _is_unstreamed_empty_message(self, item: object) -> bool: if getattr(item, "type", None) != "message" or self._message_output_index is not None: return False - return not any(getattr(part, "text", None) for part in getattr(item, "content", None) or ()) + return not any(getattr(part, "text", None) for part in attribute_items(item, "content")) + + def finished_round_event(self) -> ResponseCompletedEvent | ResponseIncompleteEvent | None: + if not isinstance(self.litellm_model_response, ModelResponse): + return None + finished: Final = self._emit_terminal_response_event(self.litellm_model_response) + if finished is not None and served_from_cache(self.litellm_custom_stream_wrapper): + finished.response.hidden_params = {**finished.response.hidden_params, "cache_hit": True} + return finished + + def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None: + self.litellm_custom_stream_wrapper.bill_failure_after_earlier_rounds(earlier) def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: """ @@ -1376,6 +1458,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped input=self.request_input, responses_api_request=dict(self.responses_api_request), + only_since_compaction=True, ), logging_obj=self.litellm_logging_obj, ) diff --git a/litellm/responses/litellm_completion_transformation/tool_search.py b/litellm/responses/litellm_completion_transformation/tool_search.py index 0a94acdd461..780a23e02b2 100644 --- a/litellm/responses/litellm_completion_transformation/tool_search.py +++ b/litellm/responses/litellm_completion_transformation/tool_search.py @@ -1,5 +1,5 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence, Set -from typing import Final, Literal, TypeAlias +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -12,6 +12,9 @@ from litellm.types.utils import LlmProviders, Message from .custom_tools import openai_shaped_tool_call_item_id +if TYPE_CHECKING: + from litellm.llms.base_llm.chat.transformation import BaseConfig + __all__ = [ "ANTHROPIC_TOOL_SEARCH_TOOL_NAME", "ANTHROPIC_TOOL_SEARCH_TOOL_TYPE", @@ -31,6 +34,7 @@ __all__ = [ "is_openai_tool_search_tool", "is_server_tool_search_tool", "provider_builds_anthropic_request", + "provider_chat_config", "runs_tool_search_on_the_server", "tool_search_call_ids", "tools_with_anthropic_tool_search_mapped", @@ -159,7 +163,7 @@ def input_without_unanswered_anthropic_tool_search_calls( return input if len(kept) == len(input) else kept -def provider_builds_anthropic_request(model: str, custom_llm_provider: str | None) -> bool: +def provider_chat_config(model: str, custom_llm_provider: str | None) -> "BaseConfig | None": from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.utils import ProviderConfigManager @@ -167,10 +171,12 @@ def provider_builds_anthropic_request(model: str, custom_llm_provider: str | Non try: llm_provider: Final = LlmProviders(provider) except ValueError: - return False - return chat_config_builds_anthropic_request( - ProviderConfigManager.get_provider_chat_config(model=model, provider=llm_provider) - ) + return None + return ProviderConfigManager.get_provider_chat_config(model=model, provider=llm_provider) + + +def provider_builds_anthropic_request(model: str, custom_llm_provider: str | None) -> bool: + return chat_config_builds_anthropic_request(provider_chat_config(model, custom_llm_provider)) def chat_config_builds_anthropic_request(chat_config: object) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index d9ebf50cc1a..a9c7818b577 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -8,7 +8,7 @@ import uuid from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial -from itertools import accumulate, chain, islice +from itertools import accumulate, chain, groupby, islice from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -29,6 +29,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( Function as NamedToolChoiceFunction, ) from openai.types.responses import ResponseFunctionToolCall, ResponseFunctionWebSearch +from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam @@ -120,6 +121,17 @@ from litellm.types.utils import ( Usage, ) +from .compaction_items import ( + carries_only_compaction, + compaction_output_item, + decode_compaction_item, + holds_compaction_summary, + provider_compaction_summary, + provider_takes_compaction_blocks, + requests_compaction, + with_answered_turns_after_newest_compaction, + without_items_before_compaction, +) from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, @@ -224,6 +236,7 @@ TOOL_CALLS_CACHE: Final = InMemoryCache() _ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_MessageT = TypeVar("_MessageT") _CACHED_CONTENT_ADAPTER: Final = TypeAdapter(ChatCompletionCachedContent) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) @@ -306,6 +319,126 @@ def _without_stored_tool_search_uses(message: SessionMessage, search_call_ids: f return message +def _instruction_messages(messages: Sequence[SessionMessage]) -> tuple[SessionMessage, ...]: + return tuple( + message + for message in messages + if isinstance(message, Mapping) and message.get("role") in ("system", "developer") + ) + + +def _message_field(message: object, key: str) -> object: + if isinstance(message, Mapping): + return _STR_KEY_DICT_ADAPTER.validate_python(message).get(key) + return getattr(message, key, None) + + +def _str_field(message: object, key: str) -> str | None: + value: Final = _message_field(message, key) + return value if isinstance(value, str) else None + + +def _tool_output_call_id(message: object) -> str | None: + return _str_field(message, "tool_call_id") if _message_field(message, "role") == "tool" else None + + +def _assistant_tool_calls(message: object) -> tuple[object, ...]: + if _message_field(message, "role") != "assistant": + return () + return tuple(_OBJECT_LIST_ADAPTER.validate_python(_message_field(message, "tool_calls") or [])) + + +def _answered_tool_call_ids(messages: Sequence[object]) -> frozenset[str]: + called: Final = frozenset( + _str_field(call, "id") for call in chain.from_iterable(map(_assistant_tool_calls, messages)) + ) + return frozenset( + call_id for call_id in map(_tool_output_call_id, messages) if call_id is not None and call_id not in called + ) + + +def _stored_turn_with_calls(message: SessionMessage, call_ids: frozenset[str]) -> SessionMessage | None: + calls: Final = [call for call in _assistant_tool_calls(message) if _str_field(call, "id") in call_ids] + if not calls: + return None + if isinstance(message, Message): + return message.model_copy( + update={"tool_calls": calls, "provider_specific_fields": _provider_fields_without_compaction(message)} + ) + return cast( # cast-ok: the stored assistant turn keeps every field but the calls the new input leaves unanswered + SessionMessage, {**_STR_KEY_DICT_ADAPTER.validate_python(message), "tool_calls": calls} + ) + + +def _call_ids_at_turn(indexed_message: tuple[int, object]) -> tuple[tuple[str | None, int], ...]: + index, message = indexed_message + return tuple((_str_field(call, "id"), index) for call in _assistant_tool_calls(message)) + + +def _newest_stored_calls_the_input_answers( + session_history: Sequence[SessionMessage], answered: frozenset[str] +) -> Iterator[SessionMessage | None]: + newest_turn_by_call_id: Final = MappingProxyType( + dict(chain.from_iterable(map(_call_ids_at_turn, enumerate(session_history)))) + ) + answered_turns: Final = tuple( + sorted((newest_turn_by_call_id[call_id], call_id) for call_id in answered if call_id in newest_turn_by_call_id) + ) + answered_by_turn: Final = MappingProxyType( + { + turn: frozenset(call_id for _, call_id in calls) + for turn, calls in groupby(answered_turns, key=lambda turn_and_call: turn_and_call[0]) + } + ) + return ( + _stored_turn_with_calls(message, answered_by_turn.get(index, frozenset())) + for index, message in enumerate(session_history) + ) + + +def with_stored_calls_the_input_answers( + new_messages: Sequence[_MessageT], session_history: Sequence[SessionMessage] +) -> tuple[_MessageT | SessionMessage, ...]: + answered: Final = _answered_tool_call_ids(new_messages) + stored_calls: Final = tuple( + turn for turn in _newest_stored_calls_the_input_answers(session_history, answered) if turn is not None + ) + if not stored_calls: + return tuple(new_messages) + first_output: Final = next( + index for index, message in enumerate(new_messages) if _tool_output_call_id(message) in answered + ) + return (*new_messages[:first_output], *stored_calls, *new_messages[first_output:]) + + +def _stored_compaction_summary(message: SessionMessage) -> str | None: + return provider_compaction_summary(message.provider_specific_fields) if isinstance(message, Message) else None + + +def _provider_fields_without_compaction(message: Message) -> Mapping[str, object] | None: + if not message.provider_specific_fields: + return None + other_fields: Final = { + key: value + for key, value in _STR_KEY_DICT_ADAPTER.validate_python(message.provider_specific_fields).items() + if key != "compaction_blocks" + } + return other_fields or None + + +def _with_stored_compaction_as_summary(message: SessionMessage) -> tuple[SessionMessage, ...]: + if not isinstance(message, Message): + return (message,) + summary: Final = provider_compaction_summary(message.provider_specific_fields) + if summary is None: + return (message,) + answer: Final = message.model_copy( + update={"provider_specific_fields": _provider_fields_without_compaction(message)} + ) + summary_message: Final = ChatCompletionUserMessage(role="user", content=summary) + return (summary_message, answer) if answer.content or answer.tool_calls else (summary_message,) + + def _input_item_cache_control(input_item: Mapping[str, object]) -> ChatCompletionCachedContent | None: cache_control: Final = input_item.get("cache_control") if not cache_control: @@ -340,6 +473,32 @@ def without_item_cache_control(input: str | ResponseInputParam) -> str | Respons ) +def _as_input_param(item: object) -> object: + return item.model_dump(exclude_none=True) if isinstance(item, ResponseCompactionItem) else item + + +def with_compaction_items_as_params(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list): + return input + return cast( # cast-ok: the same input items, a compaction item replayed from SDK output dumped to its param dict + ResponseInputParam, [_as_input_param(item) for item in input] + ) + + +def _input_to_read( + input: str | ResponseInputParam, only_since_compaction: bool, native_compaction: bool +) -> str | ResponseInputParam: + if only_since_compaction: + return without_items_before_compaction(input) + if native_compaction: + return with_answered_turns_after_newest_compaction(input) + return input + + +def _replays_compaction_block(input: str | ResponseInputParam, context_management: object) -> bool: + return requests_compaction(context_management) and holds_compaction_summary(input) + + _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai")) @@ -426,9 +585,6 @@ def _block_with_cache_control(block: object, rewrite: _CacheControlRewrite) -> o return _fields_with_cache_control(_STR_KEY_DICT_ADAPTER.validate_python(block), rewrite) -_MessageT = TypeVar("_MessageT") - - def _with_cache_control(message: _MessageT, rewrite: _CacheControlRewrite) -> _MessageT: if not isinstance(message, dict): return message @@ -870,14 +1026,21 @@ class LiteLLMCompletionResponsesConfig: ) request_kwargs: Final = _STR_KEY_DICT_ADAPTER.validate_python({**kwargs, "drop_params": drop_params}) + context_management: Final = responses_api_request.get("context_management") keeps_item_marks: Final = _target_keeps_item_marks( model, custom_llm_provider, drops_unsupported_params(request_kwargs.get("drop_params")) ) + request_input: Final = with_compaction_items_as_params(input) + native_compaction: Final = _replays_compaction_block( + request_input, context_management + ) and provider_takes_compaction_blocks(model, custom_llm_provider) bridged_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=input, + input=request_input, responses_api_request=responses_api_request, replay_reasoning=True, item_cache_control=keeps_item_marks, + native_compaction=native_compaction, + only_since_compaction=not native_compaction, ) litellm_completion_request: dict = { @@ -909,7 +1072,7 @@ class LiteLLMCompletionResponsesConfig: "response_format": response_format, "reasoning_effort": reasoning.effort, "reasoning_summary": reasoning.summary, - "context_management": responses_api_request.get("context_management"), + "context_management": context_management, # litellm specific params "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, @@ -939,6 +1102,8 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams | dict, replay_reasoning: bool = False, item_cache_control: bool = False, + native_compaction: bool = False, + only_since_compaction: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -960,6 +1125,19 @@ class LiteLLMCompletionResponsesConfig: guardrail scanning) leave both off, because they need every piece of text in the request to stay readable as message ``content``, with no item mark a strict message schema could refuse. + + ``native_compaction`` belongs to callers whose target builds an + Anthropic request that also asks for compaction in + ``context_management``: a ``compaction`` item then goes back as the + provider's own compaction block, which Anthropic rejects without that + compaction strategy. Everyone else reads its summary as a user message. + + ``only_since_compaction`` belongs to callers that send, count, or + rebuild what the model reads: the items before the newest + ``compaction`` item are dropped, keeping system and developer + messages, the way Anthropic ignores content before a compaction block. + Guardrail scans leave it off because they write rewrites back by input + item position. """ messages: list[ AllMessageValues @@ -975,10 +1153,13 @@ class LiteLLMCompletionResponsesConfig: ) ) + dumped_input: Final = with_compaction_items_as_params(input) + read_input: Final = _input_to_read(dumped_input, only_since_compaction, native_compaction) messages.extend( LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( - input=input if item_cache_control else without_item_cache_control(input), + input=read_input if item_cache_control else without_item_cache_control(read_input), replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -997,6 +1178,18 @@ class LiteLLMCompletionResponsesConfig: if (kept := _without_stored_tool_search_uses(message, search_call_ids)) is not None ) + @staticmethod + def _with_compaction_blocks_the_route_takes( + messages: Sequence[SessionMessage], route: CompletionRoute, context_management: object + ) -> tuple[SessionMessage, ...]: + if not any(_stored_compaction_summary(message) is not None for message in messages): + return tuple(messages) + if requests_compaction(context_management) and provider_takes_compaction_blocks( + route.model, route.custom_llm_provider + ): + return tuple(messages) + return tuple(chain.from_iterable(map(_with_stored_compaction_as_summary, messages))) + @staticmethod async def async_responses_api_session_handler( previous_response_id: str, @@ -1005,6 +1198,7 @@ class LiteLLMCompletionResponsesConfig: session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, request_kwargs: Mapping[str, object] = MappingProxyType({}), previous_response_client_ids: tuple[str, ...] = (), + input_replaces_history: bool = False, ) -> dict[str, object]: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1027,11 +1221,15 @@ class LiteLLMCompletionResponsesConfig: else ChatCompletionSession(messages=[], litellm_session_id=None) ) _messages: Final = litellm_completion_request.get("messages") or [] - stored_messages: Final = chat_completion_session.get("messages") or [] + route: Final = CompletionRoute(model=model, custom_llm_provider=provider) + session_history: Final = chat_completion_session.get("messages") or [] + stored_messages: Final = LiteLLMCompletionResponsesConfig._with_compaction_blocks_the_route_takes( + _instruction_messages(session_history) if input_replaces_history else session_history, + route, + request_fields.get("context_management"), + ) replayed_messages: Final = ( - LiteLLMCompletionResponsesConfig._without_tool_search_uses_off_anthropic( - stored_messages, CompletionRoute.model_validate(litellm_completion_request) - ) + LiteLLMCompletionResponsesConfig._without_tool_search_uses_off_anthropic(stored_messages, route) if any(holds_tool_search_use(message) for message in stored_messages) else stored_messages ) @@ -1062,11 +1260,16 @@ class LiteLLMCompletionResponsesConfig: # Store original _messages before combining for safety check original_new_messages: Final = _messages.copy() if _messages else [] - combined_messages = [*leading_system_messages, *session_messages, *_messages[instructions_end:]] + new_messages: Final = ( + with_stored_calls_the_input_answers(_messages[instructions_end:], session_history) + if input_replaces_history + else _messages[instructions_end:] + ) + combined_messages = [*leading_system_messages, *session_messages, *new_messages] # Fix: Ensure tool_results have corresponding tool_calls in previous assistant message # Pass tools parameter to help reconstruct tool_calls if not in cache - tools: Final = litellm_completion_request.get("tools") or [] + tools: Final = _OBJECT_LIST_ADAPTER.validate_python(litellm_completion_request.get("tools") or []) combined_messages = LiteLLMCompletionResponsesConfig._ensure_tool_results_have_corresponding_tool_calls( messages=combined_messages, tools=tools ) @@ -1115,6 +1318,7 @@ class LiteLLMCompletionResponsesConfig: def _transform_response_input_param_to_chat_completion_message( input: str | ResponseInputParam, replay_reasoning: bool = False, + native_compaction: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ]: @@ -1140,6 +1344,7 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( input_item=_input, replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -1523,7 +1728,7 @@ class LiteLLMCompletionResponsesConfig: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": - return j + return None if carries_only_compaction(messages[j]) else j return None @staticmethod @@ -1839,6 +2044,7 @@ class LiteLLMCompletionResponsesConfig: def _transform_responses_api_input_item_to_chat_completion_message( input_item: Mapping[str, object], replay_reasoning: bool = False, + native_compaction: bool = False, ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API input item into a Chat Completion message @@ -1877,6 +2083,12 @@ class LiteLLMCompletionResponsesConfig: Mapping[str, str], input_item ) ) + elif input_item.get("type") == "compaction": + return list( + LiteLLMCompletionResponsesConfig._transform_compaction_item_to_chat_completion_messages( + input_item, native_compaction + ) + ) elif input_item.get("type") == "reasoning": # A ResponseReasoningItemParam carries the prior-turn chain-of-thought. # Chat-completions providers (DeepSeek V4, Kimi K2.6, ...) expect this @@ -1928,6 +2140,23 @@ class LiteLLMCompletionResponsesConfig: return [] return [LiteLLMCompletionResponsesConfig._transform_responses_api_message_item(input_item, content)] + @staticmethod + def _transform_compaction_item_to_chat_completion_messages( + input_item: Mapping[str, object], native_compaction: bool + ) -> tuple[ChatCompletionUserMessage | ChatCompletionResponseMessage, ...]: + compaction: Final = decode_compaction_item(input_item.get("encrypted_content")) + if compaction is None: + return () + if native_compaction: + return ( + ChatCompletionResponseMessage( + role="assistant", + content=None, + provider_specific_fields={"compaction_blocks": [dict(compaction.wire_block)]}, + ), + ) + return (ChatCompletionUserMessage(role="user", content=compaction.summary),) + @staticmethod def _transform_responses_api_message_item( input_item: Mapping[str, object], content: object @@ -3038,6 +3267,7 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ]: responses_output: list[ GenericResponseOutputItem @@ -3048,8 +3278,10 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ] = [] + responses_output.extend(LiteLLMCompletionResponsesConfig._extract_compaction_output_items(choices)) responses_output.extend( LiteLLMCompletionResponsesConfig._extract_reasoning_output_items(chat_completion_response, choices) ) @@ -3115,6 +3347,11 @@ class LiteLLMCompletionResponsesConfig: output_items.append(item) return output_items + @staticmethod + def _extract_compaction_output_items(choices: Sequence[Choices]) -> tuple[ResponseCompactionItem, ...]: + item: Final = compaction_output_item(choices[0].message.provider_specific_fields) if choices else None + return () if item is None else (item,) + @staticmethod def _encode_thinking_blocks(message: Message) -> str | None: thinking_blocks: Final[Sequence[Mapping[str, object]]] = getattr(message, "thinking_blocks", None) or () @@ -3426,18 +3663,20 @@ class LiteLLMCompletionResponsesConfig: if count is not None } ) - response_usage.input_tokens_details = InputTokensDetails( - cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, - text_tokens=prompt_details.text_tokens, - audio_tokens=prompt_details.audio_tokens, - image_tokens=prompt_details.image_tokens, - video_tokens=prompt_details.video_tokens, - cached_tokens_details=( - cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None - ), - **cache_write_extra, - **cache_type_extra, - **grounding_request_counts, + response_usage.input_tokens_details = InputTokensDetails.model_validate( + { + "cached_tokens": prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, + "text_tokens": prompt_details.text_tokens, + "audio_tokens": prompt_details.audio_tokens, + "image_tokens": prompt_details.image_tokens, + "video_tokens": prompt_details.video_tokens, + "cached_tokens_details": ( + cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None + ), + **cache_write_extra, + **cache_type_extra, + **grounding_request_counts, + } ) # Translate completion_tokens_details to output_tokens_details diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 3bdb1c9c375..d993b05f5b7 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -6,10 +6,10 @@ from contextlib import contextmanager from dataclasses import dataclass from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast import httpx -from pydantic import BaseModel, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import ( TypeIs, # noqa: TID251 # narrows a client input item before its mark is read assert_never, @@ -30,6 +30,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -39,10 +40,17 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig +from litellm.responses.litellm_completion_transformation.compaction_items import ( + with_compaction_items_as_summaries, +) from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) +from litellm.responses.litellm_completion_transformation.transformation import ( + with_compaction_items_as_params, +) from litellm.responses.mcp.request_context import MCPRequestContext from litellm.responses.mcp.round_billing import ( billed_for_every_round, @@ -79,7 +87,13 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.types.responses.main import * from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypes, all_litellm_params +from litellm.types.utils import ( + ADDRESSED_RESPONSE_ID_FIELD, + CallTypes, + CostPerToken, + CustomPricingLiteLLMParams, + all_litellm_params, +) from litellm.utils import ( ProviderConfigManager, client, @@ -177,6 +191,7 @@ def mock_responses_api_response( _STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_FORWARDED_VALUE: Final = TypeAdapter(object) def _client_trace_id(request_kwargs: Mapping[str, object]) -> str | None: @@ -538,6 +553,47 @@ async def aresponses_api_with_mcp( return response +def _per_call_custom_pricing(litellm_params: GenericLiteLLMParams) -> Mapping[str, float]: + return MappingProxyType( + litellm_params.model_dump(include=set(CustomPricingLiteLLMParams.model_fields), exclude_none=True) + ) + + +class _DeploymentModelInfo(BaseModel): + model_config = ConfigDict(frozen=True) + id: object = None + + +_DEPLOYMENT_MODEL_INFO: Final[TypeAdapter[_DeploymentModelInfo | object]] = TypeAdapter( + Annotated[_DeploymentModelInfo | object, Field(union_mode="left_to_right")] +) + + +def _is_router_deployment(deployment_model_info: object) -> bool: + """A Router stamps the deployment's model_info (with its id) onto every call it routes.""" + parsed: Final = _DEPLOYMENT_MODEL_INFO.validate_python(deployment_model_info) + return isinstance(parsed, _DeploymentModelInfo) and parsed.id is not None + + +_COST_PER_TOKEN: Final = TypeAdapter(CostPerToken) + + +def _per_call_cost_per_token(litellm_params: GenericLiteLLMParams) -> CostPerToken | None: + if litellm_params.input_cost_per_token is None or litellm_params.output_cost_per_token is None: + return None + return _COST_PER_TOKEN.validate_python( + litellm_params.model_dump( + include={ + "input_cost_per_token", + "output_cost_per_token", + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + }, + exclude_none=True, + ) + ) + + def _bridges_to_chat_completions( responses_api_provider_config: BaseResponsesAPIConfig | None, use_chat_completions_api: bool ) -> bool: @@ -649,10 +705,12 @@ def _api_base_kwarg(kwargs: Mapping[str, object]) -> str | None: return api_base if isinstance(api_base, str) else None -def _dispatched_model_name(model: str, custom_llm_provider: str, api_base: str | None) -> str: - provider_model, _, _, _ = litellm.get_llm_provider( - model=model, custom_llm_provider=custom_llm_provider, api_base=api_base - ) +def _dispatched_model_name( + model: str, custom_llm_provider: str, api_base: str | None, litellm_params: GenericLiteLLMParams | None = None +) -> str: + from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata + + provider_model, _ = resolve_model_and_provider_for_metadata(model=model, custom_llm_provider=custom_llm_provider) return _strip_responses_routing_prefix(provider_model) @@ -662,6 +720,7 @@ def _will_bridge_to_chat_completions( use_chat_completions_api: bool, model_info: object, api_base: str | None, + litellm_params: GenericLiteLLMParams | None = None, ) -> bool: """``_bridges_to_chat_completions`` for callers running before the provider config is resolved. @@ -678,7 +737,7 @@ def _will_bridge_to_chat_completions( return True return _bridges_to_chat_completions( _resolve_responses_api_provider_config( - _dispatched_model_name(normalized_model[0], custom_llm_provider, api_base), + _dispatched_model_name(normalized_model[0], custom_llm_provider, api_base, litellm_params), custom_llm_provider, model_info, api_base, @@ -694,12 +753,15 @@ def provider_stores_the_conversation(model: str, custom_llm_provider: str | None first MCP round runs as an internal call that writes no spend log at all. """ api_base: Final = _api_base_kwarg(kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) return not _will_bridge_to_chat_completions( model, - custom_llm_provider or litellm.get_llm_provider(model=model, api_base=api_base)[1], + custom_llm_provider + or litellm.get_llm_provider(model=model, api_base=api_base, litellm_params=litellm_params)[1], bool(kwargs.get("use_chat_completions_api")), kwargs.get("model_info"), api_base, + litellm_params, ) @@ -811,7 +873,7 @@ async def aresponses( stream: bool | None = None, temperature: float | None = None, text: Optional["ResponseText"] = None, - text_format: type["BaseModel"] | dict | None = None, + text_format: type["BaseModel"] | dict[str, object] | None = None, tool_choice: ToolChoice | None = None, tools: Iterable[ToolParam] | None = None, top_p: float | None = None, @@ -849,7 +911,8 @@ async def aresponses( model=model, api_base=local_vars.get("base_url", None), litellm_params=GenericLiteLLMParams( - **cast("dict[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict + extra_headers=extra_headers, + **cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict ), ) # Update local_vars with detected provider (fixes #19782) @@ -880,6 +943,7 @@ async def aresponses( bool(kwargs.get("use_chat_completions_api")), kwargs.get("model_info"), _api_base_kwarg(kwargs), + GenericLiteLLMParams(**kwargs), ), ): _stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools) @@ -911,7 +975,7 @@ async def aresponses( original_model=original_model, swapped_model=model, custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs={**kwargs, "extra_headers": extra_headers}, prompt_id=prompt_id, ) kwargs.pop("prompt_id", None) @@ -995,11 +1059,14 @@ def _resolve_prompt_swapped_provider( kwargs: Mapping[str, object], prompt_id: str | None, ) -> str: - swapped_provider: Final = litellm.get_llm_provider(model=swapped_model)[1] + litellm_params: Final = GenericLiteLLMParams(**kwargs) + swapped_provider: Final = litellm.get_llm_provider(model=swapped_model, litellm_params=litellm_params)[1] if kwargs.get("api_key") is None and kwargs.get("api_base") is None: return swapped_provider try: - original_provider: Final = custom_llm_provider or litellm.get_llm_provider(model=original_model)[1] + original_provider: Final = ( + custom_llm_provider or litellm.get_llm_provider(model=original_model, litellm_params=litellm_params)[1] + ) except litellm.BadRequestError: return swapped_provider if swapped_provider == original_provider: @@ -1052,6 +1119,7 @@ def _apply_prompt_management_to_responses_call( use_chat_completions_api, kwargs.get("model_info"), _api_base_kwarg(kwargs), + GenericLiteLLMParams(**kwargs), ), ): _stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools) @@ -1086,7 +1154,7 @@ def _apply_prompt_management_to_responses_call( original_model=original_model, swapped_model=model, custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs={**kwargs, "extra_headers": local_vars.get("extra_headers")}, prompt_id=prompt_id, ) ) @@ -1300,7 +1368,7 @@ def _responses_try_dispatch_emulated_file_search( user: str | None, service_tier: str | None, safety_identifier: str | None, - text_format: type[BaseModel] | dict | None, + text_format: type[BaseModel] | dict[str, object] | None, allowed_openai_params: list[str] | None, extra_headers: dict[str, object] | None, extra_query: dict[str, object] | None, @@ -1335,7 +1403,7 @@ def _responses_try_dispatch_emulated_file_search( "background": background, "stream": stream, "temperature": temperature, - "text": text, + "text": _FORWARDED_VALUE.validate_python(text), "tool_choice": tool_choice, "top_p": top_p, "truncation": truncation, @@ -1380,7 +1448,7 @@ def responses( stream: bool | None = None, temperature: float | None = None, text: Optional["ResponseText"] = None, - text_format: type["BaseModel"] | dict | None = None, + text_format: type["BaseModel"] | dict[str, object] | None = None, tool_choice: ToolChoice | None = None, tools: Iterable[ToolParam] | None = None, top_p: float | None = None, @@ -1442,7 +1510,8 @@ def responses( model=model, api_base=local_vars.get("base_url", None), litellm_params=GenericLiteLLMParams( - **cast("dict[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict + extra_headers=extra_headers, + **cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict ), ) local_vars["custom_llm_provider"] = custom_llm_provider @@ -1458,12 +1527,26 @@ def responses( use_chat_completions_api=use_chat_completions_api, ) - # get llm provider logic - litellm_params: Final = GenericLiteLLMParams(**kwargs) + from litellm.litellm_core_utils.core_helpers import attribute_or_none + + litellm_params: Final = GenericLiteLLMParams(extra_headers=extra_headers, **kwargs) + provider_headers: Final = ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers=extra_headers, + provider_specific_header=attribute_or_none(litellm_params, "provider_specific_header"), + custom_llm_provider=custom_llm_provider, + ) + local_vars["extra_headers"] = provider_headers + + affinity_api_base: Final = ( + get_openai_credentials(api_base=litellm_params.api_base).api_base + if custom_llm_provider == "openai" + else litellm_params.api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) try: effective_extra_headers: Final = ( add_provider_affinity_header( - headers=extra_headers or MappingProxyType({}), + headers=provider_headers or MappingProxyType({}), litellm_params=MappingProxyType( { "provider_affinity_header": litellm_params.provider_affinity_header, @@ -1473,9 +1556,10 @@ def responses( "litellm_metadata": kwargs.get("litellm_metadata"), } ), + default_header=default_affinity_header, ) - if litellm_params.provider_affinity_header is not None - else extra_headers + if litellm_params.provider_affinity_header is not None or default_affinity_header is not None + else provider_headers ) except ValueError as affinity_error: raise litellm.BadRequestError( @@ -1658,6 +1742,7 @@ def responses( optional_params=dict(responses_api_request_params), litellm_params={ **responses_api_request_params, + **_per_call_custom_pricing(litellm_params), "aresponses": _is_async, "litellm_call_id": litellm_call_id, "model_info": kwargs.get("model_info"), @@ -1668,8 +1753,13 @@ def responses( custom_llm_provider=custom_llm_provider, ) + if not _is_router_deployment(deployment_model_info): + litellm_logging_obj.custom_cost_per_token = _per_call_cost_per_token(litellm_params) + litellm_logging_obj.custom_cost_per_second = litellm_params.cost_per_second # Decode any litellm-encoded encrypted-content item IDs back to their original IDs - input = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input) + input = with_compaction_items_as_summaries( + with_compaction_items_as_params(ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input)) + ) # Call the handler with _is_async flag instead of directly calling the async handler if custom_llm_provider is None: @@ -2516,7 +2606,9 @@ def compact_responses( # Decode any litellm-encoded encrypted-content item IDs back to their original IDs # before forwarding to the upstream provider. - input = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input) + input = with_compaction_items_as_summaries( + with_compaction_items_as_params(ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input)) + ) # Call the handler with _is_async flag instead of directly calling the async handler response = base_llm_http_handler.compact_response_api_handler( diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index b100858b306..cd072825ef5 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -4,6 +4,7 @@ import logging from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._internal_context import internal_sub_call @@ -24,6 +25,8 @@ from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +_DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) + class _MCPCompletionKwargs(TypedDict, total=False, extra_items=object): """Extra keywords forwarded verbatim to ``litellm.acompletion``, which owns their contract.""" @@ -566,7 +569,7 @@ async def acompletion_with_mcp( def __getattr__(self, name: str) -> object: # Delegate all other attributes to original wrapper - return getattr(self._original_wrapper, name) + return _DELEGATED_ATTRIBUTE.validate_python(getattr(self._original_wrapper, name)) # Helper class to wrap async iterator for sync iteration class _SyncIteratorWrapper: diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 73cfba8dca4..9e51e2533ea 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -3,7 +3,10 @@ from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.agentic_loop_settings import validated_max_agentic_loops +from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, @@ -287,6 +290,10 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or [] self.user_api_key_auth = user_api_key_auth self.original_request_params = original_request_params or {} + configured_rounds: Final = validated_max_agentic_loops( + self.original_request_params.get("max_agentic_loops"), field="litellm_params.max_agentic_loops" + ) + self.max_tool_call_rounds: int = MAX_MCP_TOOL_CALL_ROUNDS if configured_rounds is None else configured_rounds self.should_auto_execute = self._should_auto_execute_tools() # Streaming state management @@ -352,6 +359,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._round_max_output_index = -1 self._composed_output: list[object] = [] # mutable-ok: grows as each round finishes self._pending_mcp_call_items: list[dict[str, object]] = [] # mutable-ok: grows per executed tool + self._iteration_started = False + self._round_billing: StreamedRoundBilling | None = None + self._bridged_round_billed: LiteLLMCompletionStreamingIterator | None = None def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" @@ -436,7 +446,15 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: - chunk: Final = await self._anext_impl() + if not self._iteration_started: + self._iteration_started = True + self._round_billing = self._bill_rounds_on_one_row() + try: + chunk: Final = await self._anext_impl() + except Exception: + if self._round_billing is not None: + self._round_billing.finish() + raise sequence_number: Final = getattr(chunk, "sequence_number", None) if isinstance(sequence_number, int): if sequence_number <= self._last_sequence_number and self._last_sequence_number > 0: @@ -446,6 +464,49 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._last_sequence_number = max(self._last_sequence_number, sequence_number) return chunk + def __del__(self) -> None: + round_billing: Final[object] = getattr(self, "_round_billing", None) + if isinstance(round_billing, StreamedRoundBilling): + round_billing.release_abandoned() + + def _bill_rounds_on_one_row(self) -> StreamedRoundBilling | None: + if not self.should_auto_execute: + return None + if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator): + return StreamedRoundBilling(self.base_iterator.litellm_logging_obj) + if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator): + return StreamedRoundBilling(self.base_iterator.logging_obj) + return None + + def _restore_round_headers(self, logged_event: ResponsesAPIStreamingResponse) -> None: + if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator): + return + if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator): + self.base_iterator._restore_provider_response_headers(logged_event) + + def _round_billed_early(self) -> bool: + return self._bridged_round_billed is not None and self._bridged_round_billed is self.base_iterator + + def _bill_finished_bridged_round(self) -> None: + base: Final = self.base_iterator + if self._round_billing is None or not isinstance(base, LiteLLMCompletionStreamingIterator): + return + if self._round_billed_early(): + return + finished: Final = base.finished_round_event() + if finished is None: + return + self._bridged_round_billed = base + round_output: Final = _output_items(finished.response) + self._round_billing.complete_round( + finished, + finished.response, + round_output + if finished.type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE + else (*self._composed_output, *round_output), + self._restore_round_headers, + ) + async def _anext_impl(self) -> ResponsesAPIStreamingResponse: """ Phase-based streaming: @@ -453,7 +514,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): 2. mcp_discovery - Emit MCP discovery events (after response.output_item.added) 3. continue_initial_response - Stream the current round's response (initial or follow-up). On completion, if auto-execute is on and the response contains tool calls, loops back - through tool_execution/follow-up instead of ending, up to MAX_MCP_TOOL_CALL_ROUNDS. + through tool_execution/follow-up instead of ending, up to self.max_tool_call_rounds. 4. tool_execution - Emit tool execution events 5. finished - End iteration """ @@ -617,6 +678,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): follow-up rounds past the items already emitted, and lists every round's items on the final response.completed. """ + self._bill_finished_bridged_round() chunk_type: Final[object] = getattr(chunk, "type", None) if self._round_index > 0 and chunk_type in ( ResponsesAPIStreamEvents.RESPONSE_CREATED, @@ -630,12 +692,21 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self._output_index_offset: _set_event_field(chunk, "output_index", output_index + self._output_index_offset) + if self.should_auto_execute and chunk_type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE: + return self._charge_incomplete_round(chunk) if not (self.should_auto_execute and self._is_response_completed(chunk)): return chunk response_obj: Final[object] = getattr(chunk, "response", None) if isinstance(response_obj, ResponsesAPIResponse): self.collected_response = response_obj + if self._round_billing is not None and not self._round_billed_early(): + self._round_billing.complete_round( + chunk, + response_obj, + (*self._composed_output, *_output_items(response_obj)), + self._restore_round_headers, + ) # Move to tool execution phase after this chunk self.phase = "tool_execution" await self._generate_tool_execution_events() @@ -645,13 +716,33 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self._follow_up_pending(): self._absorb_round(response_obj) return None - if self._composed_output: - merged_output: Final[list[object]] = [ # mutable-ok: the response model declares output as a list - *self._composed_output, - *_output_items(response_obj), - ] - merged_response: Final = response_obj.model_copy(update={"output": merged_output}) - _set_event_field(chunk, "response", merged_response) + merged_response: Final = ( + response_obj.model_copy(update={"output": [*self._composed_output, *_output_items(response_obj)]}) + if self._composed_output + else response_obj + ) + client_response: Final = ( + merged_response + if self._round_billing is None + else self._round_billing.charge_final_round(chunk, merged_response, self._restore_round_headers) + ) + if client_response is not response_obj: + _set_event_field(chunk, "response", client_response) + return chunk + + def _charge_incomplete_round(self, chunk: ResponsesAPIStreamingResponse) -> ResponsesAPIStreamingResponse: + response_obj: Final[object] = getattr(chunk, "response", None) + if self._round_billing is None or not isinstance(response_obj, ResponsesAPIResponse): + return chunk + if not self._round_billed_early(): + self._round_billing.complete_round( + chunk, response_obj, _output_items(response_obj), self._restore_round_headers + ) + client_response: Final = self._round_billing.charge_final_round( + chunk, response_obj, self._restore_round_headers + ) + if client_response is not response_obj: + _set_event_field(chunk, "response", client_response) return chunk async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: @@ -929,14 +1020,14 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Remove tool_choice to avoid forcing more tool calls follow_up_params.pop("tool_choice", None) - if self.tool_call_round >= MAX_MCP_TOOL_CALL_ROUNDS: + if self.tool_call_round >= self.max_tool_call_rounds: # Hit the round cap: drop tools entirely so the model must # answer in text instead of emitting another (unexecuted) # tool call that would otherwise end the stream in silence. follow_up_params.pop("tools", None) verbose_logger.warning( - "MCP auto-execute hit MAX_MCP_TOOL_CALL_ROUNDS=%s; forcing a text-only follow-up.", - MAX_MCP_TOOL_CALL_ROUNDS, + "MCP auto-execute hit max_agentic_loops=%s; forcing a text-only follow-up.", + self.max_tool_call_rounds, ) follow_up_response: Final = await aresponses(**follow_up_params) @@ -949,6 +1040,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.collected_response = None self._cached_response_id = None self._round_index += 1 + if self._round_billing is not None: + self._round_billing.record_billed_rounds(follow_up_response) except Exception as e: verbose_logger.error("Error creating follow-up iterator: %s", e) diff --git a/litellm/responses/mcp/round_billing.py b/litellm/responses/mcp/round_billing.py index ae34ab31633..93769fbe82a 100644 --- a/litellm/responses/mcp/round_billing.py +++ b/litellm/responses/mcp/round_billing.py @@ -4,7 +4,7 @@ from contextlib import contextmanager from types import MappingProxyType from typing import Final, TypeVar -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from pydantic_core import to_jsonable_python from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params @@ -46,15 +46,18 @@ def _summed_mapping(first: Mapping[str, object], final: Mapping[str, object]) -> return MappingProxyType({key: _summed_entry(key, first, final) for key in keys}) +def _summed_usage(first: BaseModel, final: BaseModel) -> object: + return to_jsonable_python(_summed_mapping(first.model_dump(), final.model_dump()), fallback=dict) + + def merged_round_usage(first: ResponseAPIUsage | None, final: ResponseAPIUsage | None) -> ResponseAPIUsage | None: if first is None or final is None: return final if first is None else first - return ResponseAPIUsage.model_validate(_summed_mapping(first.model_dump(), final.model_dump())) + return ResponseAPIUsage.model_validate(_summed_usage(first, final)) -def _merged_chat_usage(first: Usage, final: Usage) -> Usage: - summed: Final = _summed_mapping(first.model_dump(), final.model_dump()) - return Usage.model_validate(to_jsonable_python(summed, fallback=dict)) +def merged_chat_usage(first: Usage, final: Usage) -> Usage: + return Usage.model_validate(_summed_usage(first, final)) def _chat_usage(round_response: ResponsesAPIResponse | ModelResponse) -> Usage | None: @@ -82,7 +85,7 @@ def _usage_of_both_rounds( final_usage: Final = _chat_usage(final) if first_usage is None or final_usage is None: return final_usage if first_usage is None else first_usage - return _merged_chat_usage(first_usage, final_usage) + return merged_chat_usage(first_usage, final_usage) def _repeated_fixed_margin( @@ -105,6 +108,14 @@ def _less_repeated_fixed_margin(summed: Mapping[str, object], repeated: float) - return MappingProxyType({**summed, **charged_once}) +def breakdown_less_repeated_fixed_margin( + breakdown: CostBreakdown | None, repeated_fixed_margin: float +) -> CostBreakdown | None: + if breakdown is None: + return None + return _COST_BREAKDOWN.validate_python(_less_repeated_fixed_margin(breakdown, repeated_fixed_margin)) + + def summed_cost_breakdown( first: ResponsesAPIResponse | ModelResponse, first_breakdown: CostBreakdown | None, diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py new file mode 100644 index 00000000000..443147230bf --- /dev/null +++ b/litellm/responses/mcp/stream_round_billing.py @@ -0,0 +1,376 @@ +import asyncio +from collections.abc import Callable, Coroutine, Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Final, Protocol, TypeAlias + +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse +from litellm.types.utils import CostBreakdown, ModelResponse, Usage + +Parked: TypeAlias = tuple[object, ...] +RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None] + +_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class _DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: + return logger + + +def _event_loop_running() -> bool: + try: + asyncio.get_running_loop() + except RuntimeError: + return False + return True + + +def _priced(response: ResponsesAPIResponse, price: UncachedPrice) -> ResponsesAPIResponse: + priced: Final = response.model_copy() + priced.hidden_params = {**response.hidden_params, "response_cost": price.cost} + return priced + + +def _assembled(parked: Parked | None) -> ModelResponse | None: + match parked: + case (ModelResponse() as assembled, bool()): + return assembled + case _: + return None + + +def _bridged_round_usage(parked: Parked | None) -> Usage | None: + usage: Final[object] = getattr(_assembled(parked), "usage", None) + return usage if isinstance(usage, Usage) else None + + +def _usage_of_bridged_rounds(billed: Usage | None, round_usage: Usage | None, cost: object) -> Usage | None: + if billed is None or round_usage is None: + return None + return merged_chat_usage(billed, round_usage).model_copy(update={"cost": cost}) + + +def _logged_cache_hit(response: ResponsesAPIResponse, chat_usage: Usage | None) -> bool | None: + if chat_usage is not None: + return served_from_cache(response) + return served_from_cache(response) or None + + +@dataclass(frozen=True, slots=True) +class _RoundRow: + logged: ResponsesAPIStreamingResponse | ModelResponse + cache_hit: bool | None + end_time: datetime + + +@dataclass(frozen=True, slots=True) +class _SuccessLog: + logger: LiteLLMLoggingObj + row: _RoundRow + chat_usage: Usage | None + + def carry_bridged_usage(self) -> None: + if self.chat_usage is not None: + self.logger.model_call_details["combined_usage_object"] = self.chat_usage + + +def _round_row( + completed: ResponsesAPIStreamingResponse, + response: ResponsesAPIResponse, + chat_usage: Usage | None, + restore_headers: RestoresHeaders, +) -> _RoundRow: + logged_response: Final = response.model_copy() + logged_response.hidden_params = {**response.hidden_params} + logged: Final = completed.model_copy(update={"response": logged_response}) + restore_headers(logged) + return _RoundRow( + logged, + _logged_cache_hit(response, chat_usage), + datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time + ) + + +@dataclass(frozen=True, slots=True) +class _BilledRounds: + response: ResponsesAPIResponse + breakdown: CostBreakdown | None + chat_usage: Usage | None + rounds: int = 1 + + @property + def stamped_cost(self) -> float | None: + cost: Final = self.response.hidden_params.get("response_cost") + return float(cost) if isinstance(cost, (int, float)) else None + + @property + def cost(self) -> float: + return self.stamped_cost or 0.0 + + @property + def spend(self) -> BilledRoundsSpend: + usage: Final = ( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(self.response.usage) + if self.chat_usage is None + else self.chat_usage + ) + return BilledRoundsSpend(usage, self.cost, self.breakdown) + + def with_round( + self, priced: ResponsesAPIResponse, breakdown: CostBreakdown | None, chat_usage: Usage | None + ) -> "_BilledRounds": + merged: Final = billed_for_every_round(self.response, self.breakdown, priced, breakdown) + total: Final = merged.hidden_params.get("response_cost") + charged: Final = merged.model_copy( + update={"usage": None if merged.usage is None else merged.usage.model_copy(update={"cost": total})} + ) + charged.hidden_params = { + **merged.hidden_params, + "cache_hit": served_from_cache(self.response) and served_from_cache(priced), + } + return _BilledRounds( + charged, + summed_cost_breakdown(self.response, self.breakdown, priced, breakdown), + _usage_of_bridged_rounds(self.chat_usage, chat_usage, total), + self.rounds + 1, + ) + + def row( + self, + completed: ResponsesAPIStreamingResponse, + client_output: Sequence[object], + restore_headers: RestoresHeaders, + ) -> _RoundRow: + return _round_row( + completed, + self.response.model_copy(update={"output": list(client_output)}), + self.chat_usage, + restore_headers, + ) + + def bridged_row(self, assembled: ModelResponse | None, logged_id: str) -> _RoundRow | None: + if assembled is None or self.chat_usage is None: + return None + logged: Final = assembled.model_copy(update={"id": logged_id, "usage": self.chat_usage}) + logged.hidden_params = { + **(get_hidden_params(assembled) or {}), + "response_cost": self.stamped_cost, + } + return _RoundRow( + logged, + served_from_cache(self.response), + datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time + ) + + +class StreamedRoundBilling: + """One success row per streamed auto-execute request, billed for every round it ran. + + Each round's iterator parks its success dispatch on the shared logging object rather than sending it, either as a + ready dispatch or, for a stream bridged through chat completions, as the assembled response and its cache hit. + When a round completes, before its tools run, it is folded into the rounds before it and a dispatch for every + round so far stays parked, so a client disconnect or a later round that fails still bills them. Once a later round + starts, that dispatch logs from a copy of the logging state taken then, so the later round's failure, logged on the + request's own state where the proxy reads its failure row, never reaches their success row. Only the latest + dispatch is sent, through the proxy's deferred release when the proxy armed one.""" + + def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: + self._logging_obj: Final = logging_obj + self._billed: _BilledRounds | None = None + self._banked: Parked | None = None + self._banked_row: _RoundRow | None = None + if getattr(logging_obj, _DISPATCH_RELEASE, None) is None: + setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked) + + def complete_round( + self, + completed: ResponsesAPIStreamingResponse, + response: ResponsesAPIResponse, + client_output: Sequence[object], + restore_headers: RestoresHeaders, + ) -> None: + round_parked: Final = self._take_parked() + price: Final = self._price(response, round_parked) + priced: Final = _priced(response, price) + chat_usage: Final = _bridged_round_usage(round_parked) + if self._billed is None: + first_round: Final = _BilledRounds(priced, price.cost_breakdown, chat_usage) + self._billed = first_round + self._bank( + round_parked, + first_round.bridged_row(_assembled(round_parked), priced.id) + or first_round.row(completed, client_output, restore_headers), + ) + return + billed: Final = self._billed.with_round(priced, price.cost_breakdown, chat_usage) + self._close(round_parked) + self._close(self._banked) + self._logging_obj.model_call_details.pop("combined_usage_object", None) + self._logging_obj.model_call_details.pop("response_cost", None) + self._billed = billed + self._logging_obj.cost_breakdown = billed.breakdown + row: Final = billed.bridged_row(_assembled(round_parked), billed.response.id) or billed.row( + completed, client_output, restore_headers + ) + self._bank(None if round_parked is None else (self._dispatch(row, billed, self._logging_obj),), row) + + def record_billed_rounds(self, follow_up: object) -> None: + if self._billed is None or not isinstance(follow_up, BaseResponsesAPIStreamingIterator): + return + follow_up.bill_failure_after_earlier_rounds(self._billed.spend) + self._bank_apart_from_later_rounds(self._billed) + + def charge_final_round( + self, completed: ResponsesAPIStreamingResponse, response: ResponsesAPIResponse, restore_headers: RestoresHeaders + ) -> ResponsesAPIResponse: + parked: Final = self._take_parked() + self._banked = None + if self._billed is None or self._billed.rounds == 1: + if parked is not None: + self._release(parked) + return response + self._close(parked) + charged: Final = response.model_copy(update={"usage": self._billed.response.usage}) + charged.hidden_params = {**self._billed.response.hidden_params} + if parked is not None: + final_row: Final = self._billed.bridged_row(self._latest_assembled(), charged.id) or _round_row( + completed, charged, self._billed.chat_usage, restore_headers + ) + self._release((self._dispatch(final_row, self._billed, self._logging_obj),)) + return charged + + def finish(self) -> None: + parked: Final = self._take_parked() + self._banked = None + if parked is not None: + self._release(parked) + return + self._disarm() + + def release_abandoned(self) -> None: + if not self._owns_release(): + return + if _event_loop_running(): + self.finish() + return + self.abandon() + + def abandon(self) -> None: + self._close(self._take_parked()) + self._close(self._banked) + self._banked = None + self._disarm() + + def _latest_assembled(self) -> ModelResponse | None: + row: Final = self._banked_row + return row.logged if row is not None and isinstance(row.logged, ModelResponse) else None + + def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: + if served_from_cache(response): + return _SERVED_FROM_CACHE + match round_parked: + case (ModelResponse() as assembled, bool()): + return self._logging_obj.uncached_price(assembled) + case _: + return self._logging_obj.uncached_price(response) + + def _bank(self, parked: Parked | None, row: _RoundRow) -> None: + self._banked = parked + self._banked_row = row + if parked is not None: + setattr(self._logging_obj, _PARKED_DISPATCH, parked) + + def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: + row: Final = self._banked_row + if row is None: + return + parked: Final = self._take_parked() + if parked is None: + return + self._close(parked) + self._bank((self._dispatch(row, billed, self._logging_obj.attempt_scoped_copy()),), row) + + def _dispatch( + self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj + ) -> BilledRoundsDispatch | Coroutine[object, object, None]: + dispatch: Final = self._log_success(_SuccessLog(logger, row, billed.chat_usage)) + if billed.chat_usage is None: + return dispatch + return BilledRoundsDispatch(dispatch, billed.spend) + + @staticmethod + async def _log_success(success: _SuccessLog) -> None: + success.carry_bridged_usage() + await _success_dispatcher(success.logger).dispatch_success_handlers( + success.row.logged, + start_time=success.logger.start_time, + end_time=success.row.end_time, + cache_hit=success.row.cache_hit, + prefer_async_handlers=True, + ) + + def _take_parked(self) -> Parked | None: + parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + if not isinstance(parked, tuple): + return None + setattr(self._logging_obj, _PARKED_DISPATCH, None) + return parked + + async def _run_parked(self, *parked: object) -> None: + match parked: + case (BilledRoundsDispatch(dispatch=dispatch),): + await dispatch + case (dispatch,) if asyncio.iscoroutine(dispatch): + await dispatch + case (assembled, cache_hit): + await _success_dispatcher(self._logging_obj).dispatch_success_handlers( + assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + case _: + return + + def _owns_release(self) -> bool: + release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + return release == self._run_parked + + def _release(self, parked: Parked) -> None: + if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, _PARKED_DISPATCH, parked) + return + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + spawn_logging_task( + asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging" + ) + + def _disarm(self) -> None: + if self._owns_release(): + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + + @staticmethod + def _close(parked: Parked | None) -> None: + match parked: + case (BilledRoundsDispatch(dispatch=dispatch),): + dispatch.close() + case (dispatch,) if asyncio.iscoroutine(dispatch): + dispatch.close() + case _: + return diff --git a/litellm/responses/sse_output_recovery.py b/litellm/responses/sse_output_recovery.py index cdabcb504d3..e7d0a540022 100644 --- a/litellm/responses/sse_output_recovery.py +++ b/litellm/responses/sse_output_recovery.py @@ -11,9 +11,12 @@ import json from collections.abc import Mapping from typing import Final, SupportsInt, TypeAlias, cast # noqa: TID251 # int() re-checks the cast below at runtime +from pydantic import TypeAdapter + from litellm.constants import STREAM_SSE_DONE_STRING _MAX_CONTENT_INDEX: Final = 1024 +_DECODED_JSON: Final = TypeAdapter(object) _ConvertibleToInt: TypeAlias = SupportsInt | str @@ -33,7 +36,7 @@ def parse_sse_json_chunk(chunk: str) -> dict[str, object] | None: if not stripped_chunk or stripped_chunk == STREAM_SSE_DONE_STRING or stripped_chunk.startswith("event:"): return None try: - parsed_chunk: Final[object] = json.loads(stripped_chunk) + parsed_chunk: Final = _DECODED_JSON.validate_python(json.loads(stripped_chunk)) except json.JSONDecodeError: return None if not isinstance(parsed_chunk, dict): diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c5dc45566a9..50ab3e4f75a 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -28,7 +28,7 @@ from typing import ( import httpx from openai._streaming import SSEDecoder -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import TypeIs import litellm @@ -42,6 +42,7 @@ 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 BilledRoundsSpend 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 from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -165,6 +166,7 @@ def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams: _SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache" _UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text" +_DECODED_JSON: Final = TypeAdapter(object) def _load_json_object(payload: str | bytes) -> dict[str, object]: @@ -174,7 +176,7 @@ def _load_json_object(payload: str | bytes) -> dict[str, object]: def _load_json_value(payload: str | bytes) -> object: """Parse a JSON payload whose top-level shape the caller narrows itself.""" - return json.loads(payload) + return _DECODED_JSON.validate_python(json.loads(payload)) def _model_id_from_metadata(litellm_metadata: Mapping[str, object] | None) -> str | None: @@ -363,6 +365,7 @@ class BaseResponsesAPIStreamingIterator: """ _pending_logging_tasks: tuple[asyncio.Task[object], ...] = () + _earlier_rounds_spend: BilledRoundsSpend | None = None def __init__( self, @@ -705,18 +708,22 @@ class BaseResponsesAPIStreamingIterator: if usage_obj is None: return try: - self.logging_obj.model_call_details["combined_usage_object"] = ( - ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj) - ) + failed_usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj) except (TypeError, ValueError) as usage_error: verbose_logger.debug( "could not record usage for failed responses stream: %s", usage_error, ) return - self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj.response_cost_calculator(result=response_obj) or 0.0 - ) + failed_cost: Final = self.logging_obj.response_cost_calculator(result=response_obj) or 0.0 + if self._earlier_rounds_spend is None: + self.logging_obj.model_call_details["combined_usage_object"] = failed_usage + self.logging_obj.model_call_details["response_cost"] = failed_cost + return + self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, failed_usage, failed_cost) + + def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None: + self._earlier_rounds_spend = earlier def _map_error_event_exception(self, error_obj: object) -> Exception: return _map_stream_error_to_exception( @@ -1701,6 +1708,7 @@ def _estimate_usage_from_text( messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union responses_api_request=dict(responses_api_request), + only_since_compaction=True, ) input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped model=model, messages=messages diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 503d85aaed9..2d0873a218a 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -5,14 +5,18 @@ from functools import reduce from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload from pydantic import BaseModel, TypeAdapter +from typing_extensions import TypedDict as ExtraItemsTypedDict from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire payloads without a runtime conversion import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path +from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel from litellm.types.llms.openai import ( AllMessageValues, + InputTokensDetails, OutputTokensDetails, ResponseAPIUsage, ResponseInputParam, @@ -23,13 +27,23 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import DecodedResponseId from litellm.types.utils import ( + CacheCreationTokenDetails, CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, + ProviderSpecificHeader, SpecialEnums, Usage, text_tokens_without_nested_reasoning, ) +_PROVIDER_SPECIFIC_HEADER_ADAPTER: Final[ + TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None] +] = TypeAdapter(ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None) + + +class _UsageExtraFields(ExtraItemsTypedDict, total=False, extra_items=object): + pass + def _apply_nested_drop_params(params: dict[str, object], additional_drop_params: list[str] | None) -> dict[str, object]: nested_paths: Final = tuple(path for path in additional_drop_params or () if is_nested_path(path)) @@ -41,6 +55,17 @@ def _output_token_detail(details: object, field: str) -> int | None: return value if isinstance(value, int) else None +def _cache_write_lifetimes(cache_write_tokens_details: object) -> CacheCreationTokenDetails | None: + """OpenAI Responses splits `cache_write_tokens` into `ttl_30m_tokens` and `ttl_12h_tokens`.""" + ttl_30m_tokens: Final = _output_token_detail(cache_write_tokens_details, "ttl_30m_tokens") + ttl_12h_tokens: Final = _output_token_detail(cache_write_tokens_details, "ttl_12h_tokens") + if ttl_30m_tokens is None and ttl_12h_tokens is None: + return None + return CacheCreationTokenDetails( + ephemeral_30m_input_tokens=ttl_30m_tokens, ephemeral_12h_input_tokens=ttl_12h_tokens + ) + + def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a list is a Sequence of anything return isinstance(value, list) @@ -51,6 +76,28 @@ def _is_object_dict( return isinstance(value, dict) +def _prompt_tokens_details_wrapper(details: object) -> PromptTokensDetailsWrapper: + if _is_object_dict(details): + input_details: Final = InputTokensDetails.model_validate(details) + return PromptTokensDetailsWrapper( + **input_details.model_dump(exclude={"cache_write_tokens_details"}), + cache_creation_token_details=_cache_write_lifetimes(input_details.cache_write_tokens_details), + ) + return PromptTokensDetailsWrapper( + cached_tokens=getattr(details, "cached_tokens", None), + audio_tokens=getattr(details, "audio_tokens", None), + text_tokens=getattr(details, "text_tokens", None), + image_tokens=getattr(details, "image_tokens", None), + cached_tokens_details=getattr(details, "cached_tokens_details", None), + video_tokens=getattr(details, "video_tokens", None), + cache_write_tokens=getattr(details, "cache_write_tokens", None), + cache_creation_token_details=_cache_write_lifetimes(getattr(details, "cache_write_tokens_details", None)), + cache_type=getattr(details, "cache_type", None), + web_search_requests=getattr(details, "web_search_requests", None), + google_maps_grounding_requests=getattr(details, "google_maps_grounding_requests", None), + ) + + def normalize_responses_api_stream_options( stream_options: object, ) -> ResponsesAPIStreamOptions | None: @@ -192,6 +239,20 @@ class ResponsesAPIRequestUtils: **extra_headers, } + @staticmethod + def merge_provider_specific_headers( + extra_headers: Mapping[str, object] | None, + provider_specific_header: object, + custom_llm_provider: str | None, + ) -> dict[str, object] | None: # mutable-ok: legacy Responses request fields require a dict + scoped_headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=_PROVIDER_SPECIFIC_HEADER_ADAPTER.validate_python(provider_specific_header), + custom_llm_provider=custom_llm_provider, + ) + if extra_headers is None and not scoped_headers: + return None + return {**(extra_headers or {}), **scoped_headers} + @staticmethod def _check_valid_arg( supported_params: list[str] | None, @@ -1047,7 +1108,7 @@ class ResponsesAPIRequestUtils: @staticmethod def _update_container_ids_in_response( - responses_api_response: ResponsesAPIResponse | dict[str, Any], + responses_api_response: ResponsesAPIResponse | dict[str, object], custom_llm_provider: str | None, litellm_metadata: dict[str, Any] | None = None, ) -> ResponsesAPIResponse | dict[str, object]: @@ -1061,15 +1122,16 @@ class ResponsesAPIRequestUtils: model_id: Final = model_info.get("id") # Get the output list - if isinstance(responses_api_response, dict): - output = responses_api_response.get("output", []) - else: - output = getattr(responses_api_response, "output", []) + output: Final = ( + responses_api_response.get("output", []) + if isinstance(responses_api_response, dict) + else attribute_or_none(responses_api_response, "output") + ) if not output: return responses_api_response - for item in output: + for item in iter_items(output): ResponsesAPIRequestUtils.encode_container_id_on_output_item( item=item, custom_llm_provider=custom_llm_provider, @@ -1244,27 +1306,11 @@ class ResponseAPILoggingUtils: response_api_usage = usage_input prompt_tokens: Final[int] = response_api_usage.input_tokens or 0 completion_tokens: Final[int] = response_api_usage.output_tokens or 0 - prompt_tokens_details: PromptTokensDetailsWrapper | None = None - if response_api_usage.input_tokens_details: - if isinstance(response_api_usage.input_tokens_details, dict): - prompt_tokens_details = PromptTokensDetailsWrapper(**response_api_usage.input_tokens_details) - else: - prompt_tokens_details = PromptTokensDetailsWrapper( - cached_tokens=getattr(response_api_usage.input_tokens_details, "cached_tokens", None), - audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), - text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), - image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), - cached_tokens_details=getattr( - response_api_usage.input_tokens_details, "cached_tokens_details", None - ), - video_tokens=getattr(response_api_usage.input_tokens_details, "video_tokens", None), - cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), - cache_type=getattr(response_api_usage.input_tokens_details, "cache_type", None), - web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), - google_maps_grounding_requests=getattr( - response_api_usage.input_tokens_details, "google_maps_grounding_requests", None - ), - ) + prompt_tokens_details: Final = ( + _prompt_tokens_details_wrapper(response_api_usage.input_tokens_details) + if response_api_usage.input_tokens_details + else None + ) completion_tokens_details: CompletionTokensDetailsWrapper | None = None output_tokens_details: Final[OutputTokensDetails | None] = getattr( response_api_usage, "output_tokens_details", None @@ -1288,20 +1334,22 @@ class ResponseAPILoggingUtils: audio_tokens=audio_tokens, ) - extra_usage_fields: Final = { - key: value - for key, value in (response_api_usage.model_extra or {}).items() - if key - not in ( - "input_token_details", - "output_token_details", - "prompt_tokens", - "completion_tokens", - "total_tokens", - "prompt_tokens_details", - "completion_tokens_details", - ) - } + extra_usage_fields: Final = _UsageExtraFields( + **{ + key: value + for key, value in (as_str_mapping(response_api_usage.model_extra) or {}).items() + if key + not in ( + "input_token_details", + "output_token_details", + "prompt_tokens", + "completion_tokens", + "total_tokens", + "prompt_tokens_details", + "completion_tokens_details", + ) + } + ) chat_usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, diff --git a/litellm/router.py b/litellm/router.py index a18ad7c3a97..c76e4ec843f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -38,6 +38,7 @@ from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast @@ -97,7 +98,9 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, r from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import ( declared_authenticating_provider, + inferred_provider, is_registered_custom_provider, + resolve_model_and_provider_for_metadata, ) from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR as _HIDDEN_PARAMS_ATTR, @@ -259,6 +262,7 @@ from litellm.router_utils.fallback_event_handlers import ( attempted_retries_for_request, carry_over_pre_routing_selection, carry_over_routed_deployment, + carry_superseded_attempt_usage, check_non_standard_fallback_format, clear_pre_routing_selection, committed_retry_budget_for_request, @@ -361,10 +365,11 @@ from litellm.router_utils.routing_groups import ( validate_routing_strategy, ) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch +from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage from litellm.scheduler import FlowItem, Scheduler from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.litellm_params import RoutingStrategyName +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD, RoutingStrategyName from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolParam, @@ -435,6 +440,7 @@ from litellm.utils import ( EmbeddingResponse, ModelInfoLookupMiss, ModelResponse, + ProviderConfigManager, Rules, TextCompletionResponse, function_setup, @@ -552,10 +558,12 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _CallbackT = TypeVar("_CallbackT") _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) + _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( {"model", "api_base", "api_key", "api_version", "responses_api_version"} ) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +_FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS: Final = frozenset({"chatgpt"}) _CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") _CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" @@ -563,6 +571,13 @@ CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_b _PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS: Final = frozenset({"anthropic_messages", "aresponses", "responses"}) _video_call_deployment_id: Final[ContextVar[str | None]] = ContextVar("video_call_deployment_id", default=None) +_routed_call_type: Final[ContextVar[str | None]] = ContextVar("routed_call_type", default=None) +_ROUTE_CAPABILITY_RESOLVERS: Final[Mapping[str, Callable[[str, LlmProviders], object | None]]] = MappingProxyType( + { + "adecisions": ProviderConfigManager.get_provider_decisions_config, + "decisions": ProviderConfigManager.get_provider_decisions_config, + } +) _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( { "prompt_caching": PromptCachingDeploymentCheck, @@ -819,6 +834,7 @@ def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception: _RestT: Final = TypeVar("_RestT", Iterator[object], AsyncIterator[object]) +_RouteDeploymentsT: Final = TypeVar("_RouteDeploymentsT", bound=Sequence[Mapping[str, object]] | Mapping[str, object]) @dataclass(frozen=True, slots=True) @@ -930,7 +946,7 @@ MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200 def _retry_policy_ceiling(policy: RetryPolicy) -> int: """The most retries any error class under this policy can be granted.""" - retry_counts: Final[Iterable[object]] = policy.model_dump().values() + retry_counts: Final[tuple[object, ...]] = tuple(policy.model_dump().values()) return max((retries for retries in retry_counts if isinstance(retries, int)), default=0) @@ -1046,6 +1062,8 @@ class RoutingArgs(enum.Enum): # as the per-request one built from a caller-supplied user_config, drops out on its # own rather than leaving entries behind that nothing can withdraw. _live_routers: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() +_routers_holding_global_callbacks: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() +_routers_holding_global_callbacks_lock: Final = threading.Lock() def _replay_live_router_model_cost() -> None: @@ -1232,6 +1250,7 @@ class Router: fallback_access_check: FallbackAccessCheck | None = None, fallback_budget_check: FallbackBudgetCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, + session_affinity_repin_on_failover: bool = False, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -1268,6 +1287,8 @@ class Router: alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None. provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. + session_affinity_repin_on_failover (bool): After a successful replacement response, move a session pin + whose deployment is in cooldown or removed. Async routes only. Defaults to False. Applied at startup. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. @@ -1422,6 +1443,9 @@ class Router: # Maps (team_id, team_public_model_name) -> list of indices in model_list self.team_model_to_deployment_indices: dict[tuple[str, str], list[int]] = {} self.team_public_model_names: frozenset[str] = frozenset() + # Guards model_list and the three index maps above: writers change them, and readers turn an + # index into a deployment, under it, so a lookup never lands between a removal and its re-index + self._deployment_index_lock = threading.RLock() # Initialize cache attributes that ``_invalidate_model_group_info_cache`` # and ``_invalidate_access_groups_cache`` touch *before* the first @@ -1452,6 +1476,7 @@ class Router: self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds + self.session_affinity_repin_on_failover = session_affinity_repin_on_failover self.model_group_affinity_config = model_group_affinity_config warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) @@ -1578,22 +1603,18 @@ class Router: self._init_routing_groups(self._routing_groups_input) self._override_selectors: dict[str, RouterStrategySelector | None] = {} self._override_selectors_lock = threading.Lock() + self._request_scoped_selectors: tuple[CustomLogger, ...] = () + self._request_scoped_policies: tuple[CustomLogger, ...] = () self.access_groups = None + with _routers_holding_global_callbacks_lock: + _routers_holding_global_callbacks.add(self) ## USAGE TRACKING ## - if isinstance(litellm._async_success_callback, list): - litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) - else: - litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) + litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) if isinstance(litellm.success_callback, list): litellm.logging_callback_manager.add_litellm_success_callback(self.sync_deployment_callback_on_success) else: litellm.success_callback = [self.sync_deployment_callback_on_success] - if isinstance(litellm._async_failure_callback, list): - litellm.logging_callback_manager.add_litellm_async_failure_callback( - self.async_deployment_callback_on_failure - ) - else: - litellm._async_failure_callback = [self.async_deployment_callback_on_failure] + litellm.logging_callback_manager.add_litellm_async_failure_callback(self.async_deployment_callback_on_failure) ## COOLDOWNS ## if isinstance(litellm.failure_callback, list): litellm.logging_callback_manager.add_litellm_failure_callback(self.deployment_callback_on_failure) @@ -1689,20 +1710,116 @@ class Router: # Stop contributing to cost-map rebuilds straight away rather than waiting # for this router to be collected. _live_routers.discard(self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_success_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.success_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_failure_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.failure_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.input_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.service_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, self) + self.retire_routing_state() + self._release_global_callbacks() - # Remove ForwardClientSideHeadersByModelGroup if it exists - if self.optional_callbacks is not None: - for callback in self.optional_callbacks: - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.callbacks, callback, require_self=False - ) + def _release_global_callbacks(self) -> None: + with _routers_holding_global_callbacks_lock: + _routers_holding_global_callbacks.discard(self) + slack_alerting_loggers: Final = self._slack_alerting_loggers() + adaptive_router_hooks: Final = self._adaptive_router_hooks() + vacated_keys: Final = litellm.logging_callback_manager.drop_callbacks( + ( + *self._callback_selectors(), + *self._override_selectors.values(), + *(self.optional_callbacks or ()), + *slack_alerting_loggers, + *adaptive_router_hooks, + ), + method_owners=(self, *slack_alerting_loggers), + ) + with _routers_holding_global_callbacks_lock: + survivors: Final = tuple(_routers_holding_global_callbacks) + candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors) + for candidate in litellm.logging_callback_manager.with_custom_logger_keys(candidates, vacated_keys): + self._register_router_selector(candidate) + adaptive_keys: Final = frozenset( + map(litellm.logging_callback_manager.get_custom_logger_key, adaptive_router_hooks) + ) + adaptive_survivor: Final = next((router for router in survivors if router.adaptive_routers), None) + if adaptive_survivor is not None and vacated_keys & adaptive_keys: + adaptive_survivor._sync_adaptive_router_hooks() + + def serve_through_request_logging(self) -> None: + self._request_scoped_selectors = (*self._callback_selectors(), *self._request_adaptive_router_hooks()) + self._request_scoped_policies = (*self._optional_loggers(), *self._slack_alerting_loggers()) + self._release_global_callbacks() + + def bind_to_request_logging(self, logging_obj: LiteLLMLogging) -> None: + logging_obj.add_dynamic_function_callbacks( + async_success=self.deployment_callback_on_success, + failure=self.deployment_callback_on_failure, + async_failure=self.async_deployment_callback_on_failure, + ) + for selector in self._request_scoped_selectors: + logging_obj.replace_global_callback(selector) + for policy in self._request_policies_not_set_globally(): + logging_obj.add_dynamic_callback(policy) + + def _request_policies_not_set_globally(self) -> tuple[CustomLogger, ...]: + global_keys: Final = frozenset( + map( + litellm.logging_callback_manager.get_custom_logger_key, + (callback for callback in litellm.callbacks if isinstance(callback, CustomLogger)), + ) + ) + return tuple( + policy + for policy in self._request_scoped_policies + if litellm.logging_callback_manager.get_custom_logger_key(policy) not in global_keys + ) + + def _request_adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + return tuple( + AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) + for tagged in chain.from_iterable(self.adaptive_routers.values()) + ) + + def _routing_callbacks(self) -> Iterable[object]: + if not (self._request_scoped_selectors or self._request_scoped_policies): + return litellm.callbacks + replaced_keys: Final = frozenset( + map(litellm.logging_callback_manager.get_custom_logger_key, self._request_scoped_selectors) + ) + return chain( + litellm.logging_callback_manager.without_custom_logger_keys(litellm.callbacks, replaced_keys), + self._request_scoped_selectors, + self._request_policies_not_set_globally(), + ) + + def retire_routing_state(self) -> None: + for strategy in (*self._global_callback_candidates(), *self._override_selectors.values()): + if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting): + strategy.retire() + + def _slack_alerting_loggers(self) -> tuple[CustomLogger, ...]: + slack_alerting_logger: Final = getattr(self, "slack_alerting_logger", None) + return (slack_alerting_logger,) if isinstance(slack_alerting_logger, CustomLogger) else () + + def _adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + own_strategy_ids: Final = frozenset( + id(tagged.strategy) for tagged in chain.from_iterable(self.adaptive_routers.values()) + ) + hooks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + return tuple( + hook + for hook in hooks + if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own_strategy_ids + ) + + def _optional_loggers(self) -> tuple[CustomLogger, ...]: + return tuple(callback for callback in self.optional_callbacks or () if isinstance(callback, CustomLogger)) + + def _global_callback_candidates(self) -> tuple[CustomLogger, ...]: + return (*self._callback_selectors(), *self._optional_loggers(), *self._slack_alerting_loggers()) @staticmethod def _create_redis_cache( @@ -1808,7 +1925,7 @@ class Router: return selector @staticmethod - def _register_router_selector(selector: RouterStrategySelector) -> None: + def _register_router_selector(selector: CustomLogger) -> None: if isinstance(selector, LeastBusyLoggingHandler): if isinstance(litellm.input_callback, list): litellm.logging_callback_manager.add_litellm_input_callback(selector) @@ -1817,23 +1934,28 @@ class Router: if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(selector) + def _callback_selectors(self) -> tuple[RouterStrategySelector, ...]: + return tuple( + selector + for selector in ( + *(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()), + *chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()), + ) + if selector is not None + ) + def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: """ Drop router-owned strategy selectors from litellm's global callback - lists by identity. Used before re-init (`routing_strategy_init` / - `_init_routing_groups`) so repeated `update_settings` calls don't - accumulate dead selectors that keep receiving callback events. + lists. Used before re-init (`routing_strategy_init` / + `_init_routing_groups`) and by `discard` so dead selectors stop + receiving callback events. `function_setup` copies `litellm.callbacks` + into the per-event lists on every request, so those are purged too. """ for selector in selectors: if isinstance(selector, BaseRoutingStrategy): selector.retire() - selector_ids: Final = {id(s) for s in selectors if s is not None} - if not selector_ids: - return - if isinstance(litellm.callbacks, list): - litellm.callbacks = [c for c in litellm.callbacks if id(c) not in selector_ids] - if isinstance(litellm.input_callback, list): - litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids] + litellm.logging_callback_manager.drop_callbacks(selectors) def _apply_updated_routing_strategy_args(self) -> None: """ @@ -2064,8 +2186,8 @@ class Router: public_model_name: Final = deployment.model_info.team_public_model_name if team_id is None or public_model_name is None: return False - sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or () - routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices]) + siblings: Final = self._team_deployments_named(team_id, public_model_name) or () + routable_siblings: Final = self._filter_blocked_deployments([dict(sibling) for sibling in siblings]) return len(routable_siblings) > 1 _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) @@ -2750,6 +2872,7 @@ class Router: enable_session_id_affinity=False, model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, + session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -2796,6 +2919,7 @@ class Router: enable_session_id_affinity=enable_session_id_affinity, model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, + session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -4673,15 +4797,46 @@ class Router: deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) ): kwargs.pop(key, None) - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=self.default_litellm_params) + defaults: Final = _LITELLM_PARAMS_ADAPTER.validate_python(self.default_litellm_params) + self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=defaults) model_info = deployment.get("model_info", {}).copy() deployment_litellm_model_name = deployment["litellm_params"]["model"] deployment_api_base = deployment["litellm_params"].get("api_base") deployment_model_name: Final = deployment["model_name"] - if is_clientside_credential(request_kwargs=kwargs): + from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils + + request_context: Final = cast( # cast-ok: the router's legacy request dict has string keys + Mapping[str, object], kwargs + ) + raw_deployment_params: Final[object] = deployment["litellm_params"] + deployment_context: Final = _LITELLM_PARAMS_ADAPTER.validate_python( + raw_deployment_params.model_dump(exclude_none=True) + if isinstance(raw_deployment_params, LiteLLM_Params) + else raw_deployment_params + ) + declared_provider: Final = deployment_context.get("custom_llm_provider") + declared_model: Final = deployment_context.get("model") + deployment_provider: Final = ( + declared_provider + if isinstance(declared_provider, str) + else inferred_provider(declared_model if isinstance(declared_model, str) else None) + ) + forwarded_authorization: Final = ( + ProviderSpecificHeaderUtils.forwarded_authorization( + request_context.get("provider_specific_header"), deployment_provider + ) + if deployment_provider in _FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS + else None + ) + if is_clientside_credential(request_kwargs=kwargs) or forwarded_authorization is not None: deployment_pydantic_obj: Final = self._handle_clientside_credential( - deployment=deployment, kwargs=kwargs, function_name=function_name + deployment=deployment, + kwargs={ + **kwargs, + **({"api_key": forwarded_authorization} if forwarded_authorization is not None else {}), + }, + function_name=function_name, ) model_info = deployment_pydantic_obj.model_info.model_dump() deployment_litellm_model_name = deployment_pydantic_obj.litellm_params.model @@ -4817,6 +4972,19 @@ class Router: data=_DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]), ) + deployment_litellm_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) + if router_config_requests_stream_usage( + _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + ( + _LITELLM_PARAMS_ADAPTER.validate_python( + deployment_litellm_params + if isinstance(deployment_litellm_params, Mapping) + else deployment_litellm_params.model_dump(exclude_none=True) + ), + defaults, + ), + ): + kwargs[STREAM_USAGE_INJECTED_FIELD] = True self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) if _is_fallback_hop(hop_kwargs): @@ -6076,7 +6244,7 @@ class Router: try: parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) try: - with self._video_call_pin(handler_name, kwargs): + with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, request_kwargs=kwargs, @@ -6387,7 +6555,7 @@ class Router: has_generated_content: bool, buffered_lifecycle_chunks: tuple[bytes, ...], model: str, - initial_kwargs: dict[str, Any], # mutable-ok: handed to _aanthropic_messages_fallback_attempt, which mutates it + initial_kwargs: dict[str, object], # mutable-ok: _aanthropic_messages_fallback_attempt mutates it wrapper: "FallbackAwareAnthropicMessagesStream", close_superseded: Callable[[], Awaitable[None]], ) -> AsyncGenerator[bytes, None]: @@ -6605,6 +6773,7 @@ class Router: verbose_router_logger.debug( "Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget ) + carry_superseded_attempt_usage(retry_kwargs) try: response = await attempt_function(**retry_kwargs) except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry @@ -6803,6 +6972,15 @@ class Router: finally: _video_call_deployment_id.reset(token) + @staticmethod + @contextlib.contextmanager + def _route_capability_pin(handler_name: str | None) -> Generator[None]: + token: Final = _routed_call_type.set(handler_name) + try: + yield + finally: + _routed_call_type.reset(token) + @staticmethod def _narrow_to_video_deployment(healthy_deployments: Sequence[_DeploymentT]) -> Sequence[_DeploymentT]: """ @@ -6836,7 +7014,7 @@ class Router: kwargs=kwargs, metadata_variable_name=metadata_variable_name, ) - with self._video_call_pin(handler_name, kwargs): + with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): deployment: Final = self.get_available_deployment( model=model, messages=kwargs.get("messages", None), @@ -9164,6 +9342,12 @@ class Router: if model_group is None or deployment_id is None or self.get_deployment(model_id=str(deployment_id)) is None: return + if self.session_affinity_repin_on_failover: + await self._repin_session_after_success( + _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params), + _OPAQUE_VALUE_ADAPTER.validate_python(completion_response), + ) + # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits increment_deployment_successes_for_current_minute( litellm_router_instance=self, @@ -9187,6 +9371,36 @@ class Router: "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e ) + async def _repin_session_after_success(self, litellm_params: Mapping[str, object], response: object) -> None: + if not self.session_affinity_repin_on_failover: + return + try: + from litellm.router_utils.pre_call_checks.deployment_affinity_check import completed_affinity_response + + if not completed_affinity_response(response): + return + callback: Final = next( + (logger for logger in self._optional_loggers() if isinstance(logger, DeploymentAffinityCheck)), + None, + ) + if callback is None: + return + model_group: Final = callback.get_repin_model_group(litellm_params) + if model_group is None or self._is_priority_routing_group(model_group): + return + registered_ids: Final = self.get_candidate_model_ids_for_route( + model=model_group, team_id=get_request_team_id(litellm_params) + ) + cooldown_ids: Final = frozenset( + model_id + for model_id, _ in await self.cooldown_cache.async_get_active_cooldowns( + model_ids=list(registered_ids), parent_otel_span=None + ) + ) + await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids) + except Exception as error: # noqa: BLE001 # Optional pin updates must not interrupt success accounting + verbose_router_logger.debug("Session re-pinning failed: %s", error) + async def increment_deployment_usage_for_response( self, response: object, @@ -9789,7 +10003,7 @@ class Router: Raises: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): _callback.pre_call_check(deployment) @@ -9810,7 +10024,7 @@ class Router: Raises: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): try: await _callback.async_pre_call_check(deployment, parent_otel_span) @@ -9947,7 +10161,7 @@ class Router: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ returned_healthy_deployments = healthy_deployments - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): try: returned_healthy_deployments = await _callback.async_filter_deployments( @@ -10602,6 +10816,7 @@ class Router: changes, otherwise a released router keeps recording turns through its hook.""" from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, ) for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): @@ -10611,6 +10826,8 @@ class Router: litellm.logging_callback_manager.add_litellm_callback( AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) ) + if self.adaptive_routers: + litellm.logging_callback_manager.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) def init_adaptive_router_deployment(self, deployment: Deployment) -> None: """ @@ -10646,20 +10863,17 @@ class Router: # O(k) via the name→indices map: only touch deployments whose name # is listed in `available_models`, instead of scanning model_list. for name in config.available_models: - indices = self.model_name_to_deployment_indices.get(name, []) - if not indices: + deployments = self.deployments_named(name) + if not deployments: continue - d = (self.model_list or [])[indices[0]] - mi = d.get("model_info") if isinstance(d, dict) else d.model_info - mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {}) + d = deployments[0] + mi_dict: dict[str, Any] = d.get("model_info") or {} prefs_raw = mi_dict.get("adaptive_router_preferences") if prefs_raw is not None: model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set. - lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params - lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) - cost = lp_dict.get("input_cost_per_token") + cost = d.get("litellm_params", {}).get("input_cost_per_token") if cost is None: cost = mi_dict.get("input_cost_per_token") if cost is not None: @@ -10751,10 +10965,11 @@ class Router: def set_model_list(self, model_list: list): original_model_list: Final = copy.deepcopy(model_list) self._discovered_model_info_cache.flush_cache() - self.model_list = [] - self.model_id_to_deployment_index_map = {} # Reset the index - self.model_name_to_deployment_indices = {} # Reset the model_name index - self.team_model_to_deployment_indices = {} # Reset the team_model index + with self._deployment_index_lock: + self.model_list = [] + self.model_id_to_deployment_index_map = {} # Reset the index + self.model_name_to_deployment_indices = {} # Reset the model_name index + self.team_model_to_deployment_indices = {} # Reset the team_model index self.team_pattern_routers = {} self.team_public_model_names = frozenset() # Reset per-strategy router registries so hot-reload doesn't leave @@ -10866,6 +11081,10 @@ class Router: custom_llm_provider = None dynamic_api_key = None api_base = None + elif declared_authenticating_provider(litellm_model, deployment.litellm_params.custom_llm_provider) is not None: + _model, custom_llm_provider = ( # rebind-ok: each branch resolves the deployment provider + self._model_group_llm_provider(deployment.litellm_params) + ) else: # check if model provider in supported providers ( @@ -11145,26 +11364,27 @@ class Router: - model: dict - the model to add to the list - model_id: Optional[str] - the model ID to use for indexing. If None, will try to get from model["model_info"]["id"] """ - idx: Final = len(self.model_list) - self.model_list.append(model) - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() + with self._deployment_index_lock: + idx: Final = len(self.model_list) + self.model_list.append(model) + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() - # Update model_id index for O(1) lookup - if model_id is not None: - self.model_id_to_deployment_index_map[model_id] = idx - elif model.get("model_info", {}).get("id") is not None: - self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx + # Update model_id index for O(1) lookup + if model_id is not None: + self.model_id_to_deployment_index_map[model_id] = idx + elif model.get("model_info", {}).get("id") is not None: + self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx - # Update model_name index for O(1) lookup - model_name: Final = model.get("model_name") - if model_name: - if model_name not in self.model_name_to_deployment_indices: - self.model_name_to_deployment_indices[model_name] = [] - self.model_name_to_deployment_indices[model_name].append(idx) + # Update model_name index for O(1) lookup + model_name: Final = model.get("model_name") + if model_name: + if model_name not in self.model_name_to_deployment_indices: + self.model_name_to_deployment_indices[model_name] = [] + self.model_name_to_deployment_indices[model_name].append(idx) - # Update team_model index for O(1) team-scoped lookup - self._update_team_model_index(model, idx) + # Update team_model index for O(1) team-scoped lookup + self._update_team_model_index(model, idx) def upsert_deployment(self, deployment: Deployment) -> Deployment | None: """ @@ -11191,18 +11411,15 @@ class Router: # if there is a new litellm param -> then update the deployment # remove the previous deployment - removal_idx: int | None = None - deployment_id: Final = deployment.model_info.id - deployment_fast_mapping: Final = self.model_id_to_deployment_index_map - - if deployment_id in deployment_fast_mapping: - removal_idx = deployment_fast_mapping[deployment_id] - + with self._deployment_index_lock: + removal_idx: Final = self.model_id_to_deployment_index_map.get(_deployment_model_id) if removal_idx is not None: self.model_list.pop(removal_idx) self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() - self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx) + self._update_deployment_indices_after_removal( + model_id=_deployment_model_id, removal_idx=removal_idx + ) # Free the outgoing deployment's pre-routing strategy slot (keyed by the # OLD model_name/tags) before the re-add below re-registers it. @@ -11436,35 +11653,32 @@ class Router: - The deleted deployment - OR None (if deleted deployment not found) """ - deployment_idx = None - if id in self.model_id_to_deployment_index_map: - deployment_idx = self.model_id_to_deployment_index_map[id] - try: - if deployment_idx is not None: + with self._deployment_index_lock: + deployment_idx: Final = self.model_id_to_deployment_index_map.get(id) + if deployment_idx is None: + return None # Pop the item from the list first item: Final = self.model_list.pop(deployment_idx) self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx) - _budget_limiter: Final = self.get_router_deployment_budget_limiter() - if _budget_limiter is not None: - _budget_limiter.unregister_deployment_budget(model_id=id) - try: - self._unregister_pre_routing_strategy_for_deployment( - deployment=item if isinstance(item, Deployment) else Deployment(**item) - ) - except Exception: - verbose_router_logger.exception( - "delete_deployment: could not release pre-routing strategies for model_id=%s; " - "the deployment is out of the model_list and its indices are repaired", - id, - ) - return item - else: - return None + _budget_limiter: Final = self.get_router_deployment_budget_limiter() + if _budget_limiter is not None: + _budget_limiter.unregister_deployment_budget(model_id=id) + try: + self._unregister_pre_routing_strategy_for_deployment( + deployment=item if isinstance(item, Deployment) else Deployment(**item) + ) + except Exception: + verbose_router_logger.exception( + "delete_deployment: could not release pre-routing strategies for model_id=%s; " + "the deployment is out of the model_list and its indices are repaired", + id, + ) except Exception: return None + return item def get_router_deployment_budget_limiter( self, @@ -11513,6 +11727,18 @@ class Router: if _budget_limiter is not None: _budget_limiter.register_deployment_budget(deployment=deployment.to_json(exclude_none=True)) + def deployments_named(self, model_name: str) -> tuple[DeploymentTypedDict, ...] | None: + """The deployments indexed under `model_name`, read under the index lock, or None when the name is not indexed.""" + with self._deployment_index_lock: + indices: Final = self.model_name_to_deployment_indices.get(model_name) + return None if indices is None else tuple(self.model_list[idx] for idx in indices) + + def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None: + """The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed.""" + with self._deployment_index_lock: + indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) + return None if indices is None else tuple(self.model_list[idx] for idx in indices) + def get_deployment(self, model_id: str) -> Deployment | None: """ Returns -> Deployment or None @@ -11520,9 +11746,10 @@ class Router: Raise Exception -> if model found in invalid format """ # Use O(1) lookup via model_id_to_deployment_index_map only - if model_id in self.model_id_to_deployment_index_map: - idx: Final = self.model_id_to_deployment_index_map[model_id] - model: Final = self.model_list[idx] + with self._deployment_index_lock: + idx: Final = self.model_id_to_deployment_index_map.get(model_id) + model: Final = None if idx is None else self.model_list[idx] + if model is not None: if isinstance(model, dict): return Deployment(**model) elif isinstance(model, Deployment): @@ -11556,17 +11783,17 @@ class Router: Optimized with O(1) index lookup instead of O(n) linear scan. """ # O(1) lookup in model_name index - if model_group_name in self.model_name_to_deployment_indices: - indices: Final = self.model_name_to_deployment_indices[model_group_name] - if indices: - # Return first deployment for this model_name - model: Final = self.model_list[indices[0]] - if isinstance(model, dict): - return Deployment(**model) - elif isinstance(model, Deployment): - return model - else: - raise Exception(f"Model Name invalid - {type(model)}") + with self._deployment_index_lock: + indices: Final = self.model_name_to_deployment_indices.get(model_group_name) + # Return first deployment for this model_name + model: Final = self.model_list[indices[0]] if indices else None + if model is not None: + if isinstance(model, dict): + return Deployment(**model) + elif isinstance(model, Deployment): + return model + else: + raise Exception(f"Model Name invalid - {type(model)}") return None @staticmethod @@ -11589,11 +11816,12 @@ class Router: by other teams so a shared model name never resolves another team's credentials. """ - indices: Final = self.model_name_to_deployment_indices.get(model_group_name) or () - usable: Final = ( - self.model_list[idx] for idx in indices if self.deployment_usable_by_team(self.model_list[idx], team_id) - ) - first_usable: Final = next(usable, None) + with self._deployment_index_lock: + indices: Final = self.model_name_to_deployment_indices.get(model_group_name) or () + usable: Final = ( + self.model_list[idx] for idx in indices if self.deployment_usable_by_team(self.model_list[idx], team_id) + ) + first_usable: Final = next(usable, None) if first_usable is None: return None return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable @@ -11696,11 +11924,10 @@ class Router: this never triggers pattern matching or deep copies, so it is safe to call per listed model on the /v1/models hot path. """ - indices: Final = self.model_name_to_deployment_indices.get(model_name) - if not indices: + deployments: Final = self.deployments_named(model_name) + if not deployments: return None - deployments: Final = tuple(self.model_list[index] for index in indices) model_infos: Final = tuple( MappingProxyType( { @@ -11873,10 +12100,11 @@ class Router: def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: if team_id is None: return None - team_indices: Final = self.team_model_to_deployment_indices.get((team_id, model_id)) - if not team_indices: + with self._deployment_index_lock: + team_indices: Final = self.team_model_to_deployment_indices.get((team_id, model_id)) + team_model: Final = self.model_list[team_indices[0]] if team_indices else None + if team_model is None: return None - team_model: Final = self.model_list[team_indices[0]] return Deployment(**team_model) if isinstance(team_model, dict) else team_model def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: @@ -12115,10 +12343,9 @@ class Router: Optimized with O(1) index lookup instead of O(n) linear scan. """ # O(1) lookup via model_id_to_deployment_index_map - if id in self.model_id_to_deployment_index_map: - idx: Final = self.model_id_to_deployment_index_map[id] - return self.model_list[idx] - return None + with self._deployment_index_lock: + idx: Final = self.model_id_to_deployment_index_map.get(id) + return None if idx is None else self.model_list[idx] def get_model_group(self, id: str) -> list | None: """ @@ -12211,6 +12438,17 @@ class Router: ) return model_info + @staticmethod + def _model_group_llm_provider(litellm_params: LiteLLM_Params) -> tuple[str, str]: + try: + return resolve_model_and_provider_for_metadata( + model=litellm_params.model, + custom_llm_provider=litellm_params.custom_llm_provider, + ) + except litellm.exceptions.BadRequestError as e: + verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e) + return "", "" + def _set_model_group_info( self, model_group: str, user_facing_model_group_name: str, team_id: str | None = None ) -> ModelGroupInfo | None: @@ -12299,15 +12537,7 @@ class Router: deployment_is_mapped = deployment_is_catalog_mapped(model_info, model_info_dict) # get llm provider - litellm_model, llm_provider = "", "" - try: - litellm_model, llm_provider, _, _ = litellm.get_llm_provider( - model=litellm_params.model, - custom_llm_provider=litellm_params.custom_llm_provider, - litellm_params=litellm_params, - ) - except litellm.exceptions.BadRequestError as e: - verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e) + litellm_model, llm_provider = self._model_group_llm_provider(litellm_params) if model_info is None: supported_openai_params = litellm.get_supported_openai_params( @@ -12788,20 +13018,19 @@ class Router: if model_name is not None: # O(1) lookup in model_name index, then O(k) iteration where k = deployments for this model_name - if model_name in self.model_name_to_deployment_indices: - indices: Final = self.model_name_to_deployment_indices[model_name] - for idx in indices: - model = self.model_list[idx] - if "model_info" in model and "id" in model["model_info"]: - if exclude_team_models and model["model_info"].get("team_id"): - continue - ids.append(model["model_info"]["id"]) + for model in self.deployments_named(model_name) or (): + if "model_info" in model and "id" in model["model_info"]: + if exclude_team_models and model["model_info"].get("team_id"): + continue + ids.append(model["model_info"]["id"]) else: # When model_name is None, return all model IDs # Use the index map keys for O(n) where n = total deployments - for model_id in self.model_id_to_deployment_index_map: - idx = self.model_id_to_deployment_index_map[model_id] - model = self.model_list[idx] + with self._deployment_index_lock: + id_indices: Final = tuple(self.model_id_to_deployment_index_map.items()) + models: Final = tuple(self.model_list) + for model_id, idx in id_indices: + model = models[idx] if "model_info" in model and "id" in model["model_info"]: if exclude_team_models and model["model_info"].get("team_id"): continue @@ -13079,7 +13308,9 @@ class Router: # handled downstream by the pattern_router in _common_checks_available_deployment. return None - def should_include_deployment(self, model_name: str, model: dict, team_id: str | None = None) -> bool: + def should_include_deployment( + self, model_name: str, model: DeploymentTypedDict | dict, team_id: str | None = None + ) -> bool: """ Get the team-specific model name if team_id matches the deployment. """ @@ -13127,12 +13358,10 @@ class Router: # O(1) lookup in team_model index when team_id is provided if team_id is not None: - key: Final = (team_id, model_name) - if key in self.team_model_to_deployment_indices: - indices = self.team_model_to_deployment_indices[key] + team_models: Final = self._team_deployments_named(team_id, model_name) + if team_models is not None: # O(k) where k = team deployments for this model_name (typically 1-10) - for idx in indices: - model = self.model_list[idx] + for model in team_models: if not self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): continue if model_alias is not None: @@ -13145,12 +13374,10 @@ class Router: return returned_models # O(1) lookup in model_name index - if model_name in self.model_name_to_deployment_indices: - indices = self.model_name_to_deployment_indices[model_name] - + named_models: Final = self.deployments_named(model_name) + if named_models is not None: # O(k) where k = deployments for this model_name (typically 1-10) - for idx in indices: - model = self.model_list[idx] + for model in named_models: if self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): if model_alias is not None: # Optimized: Use shallow copy since we only modify top-level model_name @@ -13164,7 +13391,7 @@ class Router: # Fallback: if team_id is provided and model_name not in index, # check if model_name matches any team_public_model_name # O(n) scan but only when team_id lookup fails - for idx, model in enumerate(self.model_list): + for _idx, model in enumerate(self.model_list): if self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): if model_alias is not None: # Optimized: Use shallow copy since we only modify top-level model_name @@ -13948,6 +14175,7 @@ class Router: input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=typed_input, responses_api_request={"instructions": instructions} if instructions is not None else {}, + only_since_compaction=True, ) return litellm.token_counter( messages=cast(list, input_messages), # cast-ok: transformed chat messages @@ -14064,6 +14292,113 @@ class Router: deployment, group, Router._params_neither_dropped_nor_allowed(call_settings, params) ) + @staticmethod + def _deployment_declares_no_vision(deployment: DeploymentTypedDict, group: str) -> bool: + plain_deployment: Final = cast(Mapping[str, object], deployment) # cast-ok: router deployments are plain dicts + model_info: Final = cast(Mapping[str, object], plain_deployment.get("model_info") or {}) # cast-ok: plain dict + declared: Final = model_info.get("supports_vision") + if declared is not None: + return declared is False + raw_params: Final = plain_deployment.get("litellm_params") or {} + params: Final = cast(Mapping[str, object], raw_params) # cast-ok: router deployments are plain dicts + if not params: + return False + base_model: Final = model_info.get("base_model") or params.get("base_model") + try: + model, custom_llm_provider = Router._capability_model_and_provider(deployment["litellm_params"], group) + return litellm.utils.is_vision_explicitly_disabled( + base_model if isinstance(base_model, str) else model, custom_llm_provider + ) + except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request + verbose_router_logger.debug( + "litellm.router.py::_deployment_declares_no_vision: keeping deployment for model=%s. Got - %s", group, e + ) + return False + + @staticmethod + def _deployment_serves_route( + deployment: DeploymentTypedDict, + group: str, + resolver: Callable[[str, LlmProviders], object | None], + ) -> bool: + if "litellm_params" not in deployment: + return True + + try: + model, resolved_custom_llm_provider = Router._capability_model_and_provider( + deployment["litellm_params"], group + ) + provider: Final = next( + (member for member in LlmProviders if member.value == resolved_custom_llm_provider), + None, + ) + return provider is not None and resolver(model, provider) is not None + except Exception as e: # noqa: BLE001 # best-effort capability filter + verbose_router_logger.debug( + "litellm.router.py::_deployment_serves_route: keeping deployment for model=%s. Got - %s", group, e + ) + return True + + @staticmethod + def _without_deployments_that_cannot_serve_route( + healthy_deployments: _RouteDeploymentsT, group: str + ) -> _RouteDeploymentsT: + handler_name: Final = _routed_call_type.get() + resolver: Final = _ROUTE_CAPABILITY_RESOLVERS.get(handler_name) if handler_name is not None else None + if resolver is None or isinstance(healthy_deployments, Mapping): + return healthy_deployments + + filtered_deployments: Final = [ + deployment + for deployment in healthy_deployments + if Router._deployment_serves_route( + cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape + group, + resolver, + ) + ] + if not filtered_deployments: + return healthy_deployments + return cast(_RouteDeploymentsT, filtered_deployments) # cast-ok: same list shape as the input candidates + + @staticmethod + def _without_declared_text_only_deployments( + healthy_deployments: _RouteDeploymentsT, + group: str, + messages: list[dict[str, str]] | None, # mutable-ok: handed to _request_carries_image as-is + request_kwargs: Mapping[str, object] | None, + ) -> _RouteDeploymentsT: + if isinstance(healthy_deployments, Mapping) or not Router._request_carries_image( + messages, request_kwargs or {} + ): + return healthy_deployments + + vision_capable_deployments: Final = [ + deployment + for deployment in healthy_deployments + if not Router._deployment_declares_no_vision( + cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape + group, + ) + ] + if not vision_capable_deployments: + return healthy_deployments + return cast(_RouteDeploymentsT, vision_capable_deployments) # cast-ok: same list shape as the input candidates + + @staticmethod + def _request_carries_image( + messages: list[dict[str, str]] | None, # mutable-ok: handed to resolve_structured_messages as-is + request_kwargs: Mapping[str, object], + ) -> bool: + from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content + from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages + + structured: Final = resolve_structured_messages( + messages=cast(list[dict[str, object]] | None, messages), # cast-ok: chat messages are plain dicts + request_kwargs=cast(dict[str, object], request_kwargs), # cast-ok: the resolver only reads the mapping + ) + return bool(structured) and request_contains_image_content(structured) + @staticmethod def _pre_call_check_requested_chat_params(request_kwargs: Mapping[str, object]) -> frozenset[str]: requested: Final = litellm.utils.get_non_default_params(passed_params=dict(request_kwargs)) @@ -14099,7 +14434,7 @@ class Router: # We only pop from the list, not modify deployment dicts - 100x+ faster on hot path (every request) _returned_deployments = list(healthy_deployments) - invalid_model_indices: Final = set() # Use set for O(1) membership checks + invalid_model_indices: Final = set[int]() # Use set for O(1) membership checks _context_window_error = False _potential_error_str = "" @@ -14321,12 +14656,13 @@ class Router: def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]: """Every team's deployments under public name `model`, for a proxy admin calling without a team.""" - team_deployments: Final = [ - self.model_list[index] - for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() - if public_model_name == model - for index in indices - ] + with self._deployment_index_lock: + team_deployments: Final = [ + self.model_list[index] + for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() + if public_model_name == model + for index in indices + ] team_ids: Final = { team_id for deployment in team_deployments @@ -14674,6 +15010,13 @@ class Router: if verbose_router_logger.isEnabledFor(logging.DEBUG): verbose_router_logger.debug("healthy_deployments after web search filter: %s", healthy_deployments) + healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows + self._without_deployments_that_cannot_serve_route( + healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type + model, + ) + ) + if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: raise litellm.ServiceUnavailableError( @@ -14688,10 +15031,12 @@ class Router: ) # Health-check-based filtering (before cooldown) - healthy_deployments = await self._async_filter_health_check_unhealthy_deployments( - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - health_check_probe=health_check_probe, + healthy_deployments = ( # rebind-ok: filtered candidates enter health checks + await self._async_filter_health_check_unhealthy_deployments( + healthy_deployments=healthy_deployments, + parent_otel_span=parent_otel_span, + health_check_probe=health_check_probe, + ) ) routing_read_batch: Final = RoutingReadBatch.active() @@ -14765,8 +15110,11 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts - list[dict[str, object]], healthy_deployments + batch_candidates: Final = self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( batch_candidates, @@ -15942,6 +16290,12 @@ class Router: specific_deployment=specific_deployment, request_kwargs=request_kwargs, ) + healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows + self._without_deployments_that_cannot_serve_route( + healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type + model, + ) + ) strategy, strategy_selector = self._get_routing_context(model, request_kwargs) routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs) @@ -15999,8 +16353,11 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts - list[dict[str, object]], healthy_deployments + batch_candidates: Final = self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( batch_candidates, diff --git a/litellm/router_strategy/adaptive_router/__init__.py b/litellm/router_strategy/adaptive_router/__init__.py index d7f55ebced9..1f7e6b93169 100644 --- a/litellm/router_strategy/adaptive_router/__init__.py +++ b/litellm/router_strategy/adaptive_router/__init__.py @@ -1,6 +1,9 @@ """Adaptive router strategy. See README.md for design overview.""" from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter -from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, +) -__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"] +__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook", "AdaptiveRouterResponseHeaderHook"] diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 92207e07c52..3427b81be68 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -170,12 +170,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequen return content, tool_calls -class AdaptiveRouterPostCallHook(CustomLogger): - """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" - - def __init__(self, adaptive_router: AdaptiveRouter) -> None: - self.adaptive_router = adaptive_router - +class AdaptiveRouterResponseHeaderHook(CustomLogger): async def async_post_call_response_headers_hook( self, data: Mapping[str, object], @@ -200,6 +195,13 @@ class AdaptiveRouterPostCallHook(CustomLogger): return None return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} + +class AdaptiveRouterPostCallHook(CustomLogger): + """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" + + def __init__(self, adaptive_router: AdaptiveRouter) -> None: + self.adaptive_router = adaptive_router + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): await self._record(kwargs, response_obj, response_status=200) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 51400e6da24..2cea7389808 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -24,6 +24,8 @@ class BaseRoutingStrategy(ABC): self.dual_cache = dual_cache self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] self._sync_task: asyncio.Task[None] | None = None + self._retired = False + self._flush_after_retire: asyncio.Task[None] | None = None if should_batch_redis_writes: self.setup_sync_task(default_sync_interval) @@ -47,13 +49,25 @@ class BaseRoutingStrategy(ABC): def retire(self) -> None: self.cancel_sync_task() - if not self.redis_increment_operation_queue: + self._retired = True + self._schedule_flush_after_retire() + + def _schedule_flush_after_retire(self) -> None: + if self._flush_after_retire is not None and not self._flush_after_retire.done(): + return + if not self.redis_increment_operation_queue or self.dual_cache.redis_cache is None: return try: loop: Final = asyncio.get_running_loop() except RuntimeError: return - loop.create_task(self._push_in_memory_increments_to_redis()) + self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) + + async def _drain_increments_after_retire(self) -> None: + if not self.redis_increment_operation_queue: + return + await self._push_in_memory_increments_to_redis() + await self._drain_increments_after_retire() async def cleanup(self): """Cleanup method to be called when shutting down""" @@ -98,6 +112,8 @@ class BaseRoutingStrategy(ABC): self.redis_increment_operation_queue.append(increment_op) self.add_to_in_memory_keys_to_update(key=key) + if self._retired: + self._schedule_flush_after_retire() return result @with_service_target("router_usage") diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 147ce60a99b..eab232eae18 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -128,7 +128,9 @@ class RouterBudgetLimiting(CustomLogger): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None - asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._sync_task: asyncio.Task[None] = asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._retired = False + self._flush_after_retire: asyncio.Task[None] | None = None self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config self.deployment_budget_config: GenericBudgetConfigType | None = None self.tag_budget_config: GenericBudgetConfigType | None = None @@ -140,6 +142,30 @@ class RouterBudgetLimiting(CustomLogger): if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(self) + def retire(self) -> None: + self._sync_task.cancel() + self._retired = True + self._schedule_flush_after_retire() + + def _schedule_flush_after_retire(self) -> None: + if self._flush_after_retire is not None and not self._flush_after_retire.done(): + return + if not self._has_unwritten_increments(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) + + def _has_unwritten_increments(self) -> bool: + return bool(self.redis_increment_operation_queue) or self._detached_increment_operations is not None + + async def _drain_increments_after_retire(self) -> None: + while self._has_unwritten_increments(): + if not await self._push_in_memory_increments_to_redis(): + await asyncio.sleep(DEFAULT_REDIS_SYNC_INTERVAL) + @with_service_target("router_budgets") async def async_filter_deployments( self, @@ -427,6 +453,8 @@ class RouterBudgetLimiting(CustomLogger): if self.dual_cache.redis_cache is None: return self.redis_increment_operation_queue.append(increment_op) + if self._retired: + self._schedule_flush_after_retire() def _get_redis_increment_queue_lock(self) -> asyncio.Lock: return self._redis_increment_queue_lock diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a475358e1f0..cc592bbcbc5 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -3126,17 +3126,14 @@ class ComplexityRouter(CustomLogger): model_to_prefs: Final[dict[str, AdaptiveRouterPreferences]] = {} model_to_cost: Final[dict[str, float]] = {} - model_list: Final = getattr(self.litellm_router_instance, "model_list", None) or [] - name_to_indices: Final = getattr(self.litellm_router_instance, "model_name_to_deployment_indices", {}) or {} for name in available_models: - indices = name_to_indices.get(name, []) - if not indices: + deployments = self.litellm_router_instance.deployments_named(name) + if not deployments: model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) model_to_cost[name] = 0.0 continue - deployment = model_list[indices[0]] - mi = deployment.get("model_info") if isinstance(deployment, dict) else deployment.model_info - mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {}) + deployment = deployments[0] + mi_dict: dict[str, Any] = deployment.get("model_info") or {} prefs_raw = mi_dict.get("adaptive_router_preferences") if prefs_raw is not None: model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) @@ -3144,9 +3141,7 @@ class ComplexityRouter(CustomLogger): model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set. - lp = deployment.get("litellm_params") if isinstance(deployment, dict) else deployment.litellm_params - lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) - cost = lp_dict.get("input_cost_per_token") + cost = deployment.get("litellm_params", {}).get("input_cost_per_token") if cost is None: cost = mi_dict.get("input_cost_per_token") model_to_cost[name] = float(cost) if cost is not None else 0.0 diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 4a2c56daf00..ec960a19f79 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -7,12 +7,12 @@ All values are configurable via proxy config.yaml. import math import re -import warnings +import sys from collections.abc import Iterable, Mapping from dataclasses import dataclass from enum import Enum from types import MappingProxyType -from typing import Annotated, Final, Literal, NamedTuple +from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple from pydantic import ( ConfigDict, @@ -27,8 +27,13 @@ from pydantic import ( from litellm.types.llms.base import LiteLLMBaseModel -with warnings.catch_warnings(): - warnings.simplefilter("ignore", DeprecationWarning) +if TYPE_CHECKING: + import sre_constants + import sre_parse +elif sys.version_info >= (3, 11): + from re import _constants as sre_constants + from re import _parser as sre_parse +else: import sre_constants import sre_parse diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 83bb22435e6..4e994d3fdc6 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -3,7 +3,7 @@ import random from collections.abc import Sequence from datetime import datetime, timedelta -from math import ceil +from math import ceil, isfinite from typing import TYPE_CHECKING, Any, Final from pydantic import Field @@ -14,7 +14,9 @@ from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, safe_divide_seconds +from litellm.responses.utils import ResponseAPILoggingUtils from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import LiteLLMPydanticObjectBase if TYPE_CHECKING: @@ -44,6 +46,16 @@ def _percentile_latency(samples: Sequence[float], percentile: float) -> float: return values[index] +def _usable_latency_samples(samples: Sequence[float]) -> tuple[float, ...]: + """Drop samples that cannot be a real elapsed time. + + ``completion_start_time - start_time`` goes negative when the timestamps invert. + A negative or NaN sample pushes ``lowest_latency + buffer`` below every candidate, + leaving the deployment filter empty. + """ + return tuple(sample for sample in samples if isfinite(sample) and sample >= 0) + + def _ttft_seconds(elapsed: timedelta | float) -> float: if isinstance(elapsed, timedelta): return elapsed.total_seconds() @@ -86,7 +98,7 @@ class LowestLatencyLoggingHandler(CustomLogger): { {model_group}_map: { id: { - "latency": [..] + "latency_normalized": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } @@ -116,11 +128,16 @@ class LowestLatencyLoggingHandler(CustomLogger): final_value: float = response_ms total_tokens = 0 - if isinstance(response_obj, ModelResponse): - _usage: Final = getattr(response_obj, "usage", None) + if isinstance(response_obj, (ModelResponse, ResponsesAPIResponse)): + _usage: Final = ( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response_obj.usage) + if isinstance(response_obj, ResponsesAPIResponse) + else getattr(response_obj, "usage", None) + ) if _usage is not None: completion_tokens: Final = _usage.completion_tokens - total_tokens = _usage.total_tokens + if isinstance(response_obj, ModelResponse): + total_tokens = _usage.total_tokens # response_ms is already normalized to float seconds above response_seconds: Final = response_ms @@ -143,10 +160,12 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: - request_count_dict[id].setdefault("latency", []).append(final_value) + if len(request_count_dict[id].get("latency_normalized", [])) < self.routing_args.max_latency_list_size: + request_count_dict[id].setdefault("latency_normalized", []).append(final_value) else: - request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [final_value] + request_count_dict[id]["latency_normalized"] = request_count_dict[id]["latency_normalized"][1:] + [ + final_value + ] ## Time to first token if time_to_first_token is not None: @@ -212,7 +231,7 @@ class LowestLatencyLoggingHandler(CustomLogger): { {model_group}_map: { id: { - "latency": [..] + "latency_normalized": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } @@ -225,10 +244,15 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - give 1000s penalty for failing - if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: - request_count_dict[id].setdefault("latency", []).append(1000.0) + if ( + len(request_count_dict[id].get("latency_normalized", [])) + < self.routing_args.max_latency_list_size + ): + request_count_dict[id].setdefault("latency_normalized", []).append(1000.0) else: - request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [1000.0] + request_count_dict[id]["latency_normalized"] = request_count_dict[id]["latency_normalized"][ + 1: + ] + [1000.0] await self.router_cache.async_set_cache( key=latency_key, @@ -270,7 +294,7 @@ class LowestLatencyLoggingHandler(CustomLogger): { {model_group}_map: { id: { - "latency": [..] + "latency_normalized": [..] "time_to_first_token_seconds": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } @@ -300,11 +324,16 @@ class LowestLatencyLoggingHandler(CustomLogger): final_value: float = response_ms total_tokens = 0 - if isinstance(response_obj, ModelResponse): - _usage: Final = getattr(response_obj, "usage", None) + if isinstance(response_obj, (ModelResponse, ResponsesAPIResponse)): + _usage: Final = ( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response_obj.usage) + if isinstance(response_obj, ResponsesAPIResponse) + else getattr(response_obj, "usage", None) + ) if _usage is not None: completion_tokens: Final = _usage.completion_tokens - total_tokens = _usage.total_tokens + if isinstance(response_obj, ModelResponse): + total_tokens = _usage.total_tokens # response_ms is already normalized to float seconds above response_seconds: Final = response_ms @@ -331,10 +360,12 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: - request_count_dict[id].setdefault("latency", []).append(final_value) + if len(request_count_dict[id].get("latency_normalized", [])) < self.routing_args.max_latency_list_size: + request_count_dict[id].setdefault("latency_normalized", []).append(final_value) else: - request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [final_value] + request_count_dict[id]["latency_normalized"] = request_count_dict[id]["latency_normalized"][1:] + [ + final_value + ] ## Time to first token if time_to_first_token is not None: @@ -403,7 +434,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ## if healthy deployment not yet used if d["model_info"]["id"] not in all_deployments: all_deployments[d["model_info"]["id"]] = { - "latency": [0], + "latency_normalized": [0], precise_minute: {"tpm": 0, "rpm": 0}, } @@ -443,21 +474,22 @@ class LowestLatencyLoggingHandler(CustomLogger): or _deployment.get("model_info", {}).get("rpm", None) or float("inf") ) - item_latency = item_map.get("latency", []) + item_latency = item_map.get("latency_normalized", []) item_ttft_latency = item_map.get("time_to_first_token_seconds", []) item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) + ttft_samples = _usable_latency_samples(item_ttft_latency) use_ttft = ( request_kwargs is not None and request_kwargs.get("stream", None) is not None and request_kwargs["stream"] is True - and len(item_ttft_latency) > 0 + and len(ttft_samples) > 0 ) selected_latency = ( - _percentile_latency(item_ttft_latency, self.routing_args.ttft_percentile) + _percentile_latency(ttft_samples, self.routing_args.ttft_percentile) if use_ttft and self.routing_args.ttft_percentile is not None - else _average_latency(item_ttft_latency if use_ttft else item_latency) + else _average_latency(ttft_samples if use_ttft else _usable_latency_samples(item_latency)) ) # -------------- # @@ -491,10 +523,9 @@ class LowestLatencyLoggingHandler(CustomLogger): # Find deployments within buffer of lowest latency buffer: Final = self.routing_args.lowest_latency_buffer * lowest_latency - valid_deployments: Final = [x for x in sorted_deployments if x[1] <= lowest_latency + buffer] + valid_deployments: Final = tuple(x for x in sorted_deployments if x[1] <= lowest_latency + buffer) - # Pick a random deployment from valid deployments - random_valid_deployment: Final = random.choice(valid_deployments) + random_valid_deployment: Final = random.choice(valid_deployments or sorted_deployments) deployment = random_valid_deployment[0] metadata_field: Final = self._select_metadata_field(request_kwargs) if request_kwargs is not None and metadata_field in request_kwargs: diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index 1d119375778..d8c52b4c8b6 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -47,13 +47,13 @@ class QualityRouter(CustomLogger): model_name: str, litellm_router_instance: "Router", default_model: str | None = None, - quality_router_config: dict[str, Any] | None = None, + quality_router_config: Mapping[str, object] | None = None, ): self.model_name = model_name self.litellm_router_instance = litellm_router_instance if quality_router_config: - self.config = QualityRouterConfig(**quality_router_config) + self.config = QualityRouterConfig.model_validate(dict(**quality_router_config)) else: self.config = QualityRouterConfig() @@ -185,7 +185,7 @@ class QualityRouter(CustomLogger): if isinstance(raw_prefs, RoutingPreferences): prefs = raw_prefs elif isinstance(raw_prefs, dict): - prefs = RoutingPreferences(**raw_prefs) + prefs = RoutingPreferences.model_validate(dict(**raw_prefs)) else: # A Pydantic object of some other shape — coerce via its dict. prefs = RoutingPreferences( diff --git a/litellm/router_strategy/savings_baseline.py b/litellm/router_strategy/savings_baseline.py index b5b80ddae6d..add0f78fd3a 100644 --- a/litellm/router_strategy/savings_baseline.py +++ b/litellm/router_strategy/savings_baseline.py @@ -14,7 +14,7 @@ bare string with no provider beside them; an operator who writes ``deepseek-r1`` Azure would otherwise be priced against whoever else owns that name. """ -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Final, NamedTuple from litellm._logging import verbose_router_logger @@ -80,12 +80,11 @@ def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]: straight at a provider model rather than at a configured group, and prices under its own name because there is no deployment to override it. """ - indices: Final = router.model_name_to_deployment_indices.get(group_name) - if not indices: + deployments: Final = router.deployments_named(group_name) + if not deployments: return (Baseline(qualified),) if (qualified := canonical_model(group_name)) else () - def candidate(index: int) -> Baseline | None: - deployment: Final = router.model_list[index] + def candidate(deployment: Mapping[str, object]) -> Baseline | None: params: Final = deployment.get("litellm_params") if not isinstance(params, dict): return None @@ -98,7 +97,7 @@ def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]: deployment_id: Final = info.get("id") if isinstance(info, dict) else None return Baseline(qualified, str(deployment_id) if deployment_id else None) - return tuple(c for index in indices if (c := candidate(index)) is not None) + return tuple(c for deployment in deployments if (c := candidate(deployment)) is not None) def _priced(router: "Router", candidate: Baseline) -> tuple[float, Baseline] | None: diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index e1201df1a13..6c93fe6d239 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -8,6 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_router_logger @@ -15,6 +16,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.litellm_logging import Logging, same_attempt_marker from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, @@ -51,6 +53,7 @@ else: # (e.g. a nonexistent file/batch/thread id), independent of the selected deployment's health. _REQUEST_SCOPED_STATUS_CODES: Final = frozenset((404,)) _NO_MODEL_CALL_DETAILS: Final[Mapping[str, object]] = MappingProxyType({}) +_FALLBACK_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) def _trigger_cooldown_for_failed_deployment( @@ -447,6 +450,24 @@ def mid_stream_retry_kwargs( return {key: value for key, value in hop_kwargs.items() if key not in _MID_STREAM_RETRY_STRIPPED_KEYS} +def carry_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None: + """Carry the superseded stream's partial usage into billing if the request ultimately fails.""" + logging_obj: Final = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.carry_failed_attempt_usage() + + +def _failure_callbacks_blocked_logging_object(kwargs: Mapping[str, object]) -> Logging | None: + logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return None + if logging_obj.model_call_details.get("has_logged_async_failure", False) is not True: + return None + current_marker: Final = logging_obj.model_call_details.get("current_attempt_api_call_start_time") + logged_marker: Final = logging_obj.model_call_details.get("async_failure_logged_attempt") + return logging_obj if not same_attempt_marker(current_marker, logged_marker) else None + + def _request_metadata_bucket(kwargs: Mapping[str, object]) -> Mapping[str, object] | None: bucket: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) return bucket if isinstance(bucket, Mapping) else None @@ -869,6 +890,7 @@ async def run_async_fallback( kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors + carry_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) verbose_router_logger.info("Successful fallback b/w models.") response = add_fallback_headers_to_response( @@ -891,13 +913,16 @@ async def run_async_fallback( kwargs=kwargs, original_exception=original_exception, ) - logging_obj = kwargs.get("litellm_logging_obj") - if logging_obj is not None and logging_obj.model_call_details.get("has_logged_async_failure", False): + if ( + failure_callbacks_blocked_logging_obj := _failure_callbacks_blocked_logging_object( + _FALLBACK_KWARGS_ADAPTER.validate_python(kwargs) + ) + ) is not None: _trigger_cooldown_for_failed_deployment( litellm_router=litellm_router, kwargs=kwargs, exception=e, - model_call_details=logging_obj.model_call_details, + model_call_details=failure_callbacks_blocked_logging_obj.model_call_details, ) raise error_from_fallbacks diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 88fb8876175..8693eaa5151 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -16,6 +16,7 @@ import hashlib from collections.abc import Callable, Mapping, Sequence from typing import Final, cast +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target @@ -38,6 +39,27 @@ class DeploymentAffinityCacheValue(TypedDict, closed=True): model_id: ReadOnly[str] +def _affinity_response_fields(response: object) -> Mapping[str, object]: + return ( + TypeAdapter(Mapping[str, object]).validate_python(response) + if isinstance(response, Mapping) + else {name: getattr(response, name, None) for name in ("response", "status", "error")} + ) + + +def completed_affinity_response(response: object) -> bool: + if response is None: + return False + envelope: Final = _affinity_response_fields(response) + if envelope.get("error") is not None or envelope.get("status") not in (None, "completed"): + return False + nested: Final = envelope.get("response") + fields: Final = _affinity_response_fields(nested) if nested is not None else envelope + return ( + fields.get("response") is None and fields.get("error") is None and fields.get("status") in (None, "completed") + ) + + VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( { "deployment_affinity", @@ -86,6 +108,7 @@ class DeploymentAffinityCheck(CustomLogger): enable_session_id_affinity: bool = False, model_group_affinity_config: dict[str, list[str]] | None = None, is_priority_group: Callable[[str], bool] | None = None, + session_affinity_repin_on_failover: bool = False, ): super().__init__() self.cache = cache @@ -95,6 +118,7 @@ class DeploymentAffinityCheck(CustomLogger): self.enable_session_id_affinity = enable_session_id_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} self.is_priority_group = is_priority_group + self.session_affinity_repin_on_failover = session_affinity_repin_on_failover def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -297,7 +321,7 @@ class DeploymentAffinityCheck(CustomLogger): return f"{cls.USER_ID_AFFINITY_PREFIX}{user_id}" @staticmethod - def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: + def _get_session_id_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): session_id = DeploymentAffinityCheck._get_session_id_from_metadata_dict(metadata=metadata) if session_id is not None: @@ -305,7 +329,7 @@ class DeploymentAffinityCheck(CustomLogger): return None @staticmethod - def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None: + def _get_marker_session_affinity_ttl(request_kwargs: Mapping[str, object]) -> int | None: """TTL from the session-affinity marker the Router stamps at pre-routing time when an auto-router routed this request with session_affinity enabled. Marker presence enables session pinning for this request only; anything that @@ -357,7 +381,7 @@ class DeploymentAffinityCheck(CustomLogger): model: str, healthy_deployments: list, messages: list[AllMessageValues] | None, - request_kwargs: dict | None = None, + request_kwargs: dict[str, object] | None = None, parent_otel_span: Span | None = None, ) -> list[dict]: """ @@ -418,7 +442,15 @@ class DeploymentAffinityCheck(CustomLogger): session_cache_key: Final = self.get_session_affinity_cache_key( model_group=stable_model_map_key, session_id=session_id, user_key=user_key ) - session_cache_result: Final = await self.cache.async_get_cache(key=session_cache_key) + session_cache_result: Final = ( + await self._read_shared_session_pin( + session_cache_key, + self._get_marker_session_affinity_ttl(request_kwargs) or self.ttl_seconds, + parent_otel_span, + ) + if self.session_affinity_repin_on_failover + else await self.cache.async_get_cache(key=session_cache_key) + ) session_model_id: str | None = None if isinstance(session_cache_result, dict): @@ -482,6 +514,71 @@ class DeploymentAffinityCheck(CustomLogger): ) return [deployment] + async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: + redis_cache: Final = self.cache.redis_cache + if redis_cache is None: + return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key)) + local_before: Final = TypeAdapter(object).validate_python(self.cache.in_memory_cache.get_cache(cache_key)) + try: + shared: Final[object] = await redis_cache.async_get_cache( + key=cache_key, parent_otel_span=parent_otel_span, raise_on_error=True + ) + except Exception: # noqa: BLE001 # Optional affinity falls back locally on Redis admission or read failure + verbose_router_logger.debug("Session pin read failed, using pod-local affinity") + return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key, local_only=True)) + if self.cache.in_memory_cache.get_cache(cache_key) is local_before: + if shared is None: + self.cache.in_memory_cache.delete_cache(cache_key) + else: + self._set_local_pin(cache_key, shared, ttl_seconds) + return shared + + def get_repin_model_group(self, request_kwargs: Mapping[str, object]) -> str | None: + model_group: Final = next( + ( + value + for metadata in self._iter_metadata_dicts(request_kwargs) + if isinstance(value := metadata.get("deployment_model_name"), str) and value + ), + None, + ) + if model_group is None or self._get_session_id_from_request_kwargs(request_kwargs) is None: + return None + if ( + not self._get_effective_flags(model_group)[2] + and self._get_marker_session_affinity_ttl(request_kwargs) is None + ): + return None + return model_group + + async def async_repin_session_after_success( + self, + request_kwargs: Mapping[str, object], + model_group: str, + eligible_deployment_ids: frozenset[str], + ) -> None: + marker_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs) + if not self._get_effective_flags(model_group)[2] and marker_ttl is None: + return + session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) + if session_id is None: + return + model_info: Final = TypeAdapter(Mapping[str, object]).validate_python(request_kwargs.get("model_info") or {}) + model_id: Final = model_info.get("id") + if model_id is None or str(model_id) not in eligible_deployment_ids: + return + cache_key: Final = self.get_session_affinity_cache_key( + model_group, session_id, self.get_user_key_from_request_kwargs(request_kwargs) + ) + await claim_affinity_pin( + self.cache, + cache_key, + DeploymentAffinityCacheValue(model_id=str(model_id)), + marker_ttl if marker_ttl is not None else self.ttl_seconds, + eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_deployment_ids), + preserve_legacy_model_ids=True, + ) + async def async_pre_call_deployment_hook( self, kwargs: dict[str, object], call_type: CallTypes | None ) -> dict | None: diff --git a/litellm/router_utils/stream_usage_injection.py b/litellm/router_utils/stream_usage_injection.py new file mode 100644 index 00000000000..946543d9b83 --- /dev/null +++ b/litellm/router_utils/stream_usage_injection.py @@ -0,0 +1,17 @@ +from collections.abc import Mapping, Sequence + + +def router_config_requests_stream_usage( + request_kwargs: Mapping[str, object], configured_params: Sequence[Mapping[str, object]] +) -> bool: + if request_kwargs.get("stream") is not True or "stream_options" in request_kwargs: + return False + return any(_includes_usage(params.get("stream_options")) for params in configured_params) + + +def _includes_usage(stream_options: object) -> bool: + match stream_options: + case {"include_usage": True}: + return True + case _: + return False diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index ebb221e98be..74cd28161a4 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -165,6 +165,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { } MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { + Feature.PROVIDER_AFFINITY, Feature.SAFEGUARDS_CLASSIFIER, Feature.PRE_REQUEST_HOOK, Feature.WEBSEARCH_INTERCEPTION, diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index 0ad1b787e38..0af3a2e6fa6 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -128,7 +128,7 @@ def _custom_loggers(callbacks: Iterable[object]) -> tuple[CustomLogger, ...]: def overridden(method: str, callbacks: Iterable[object]) -> bool: - base: Final = getattr(CustomLogger, method) + base: Final[object] = getattr(CustomLogger, method) return any(getattr(type(callback), method, base) is not base for callback in _custom_loggers(callbacks)) diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 691c04f2daa..b53e80805a5 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -5,10 +5,15 @@ from __future__ import annotations from typing import Final import litellm -from litellm.rust_bridge.detectors import overridden, present, streaming +from litellm.llms.opencode.common_utils import is_opencode_api_base +from litellm.rust_bridge.detectors import api_base, overridden, present, streaming from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str +def provider_affinity(request: NativeCall) -> bool: + return is_opencode_api_base(api_base(request)) + + def anthropic_workload_identity(request: NativeCall) -> bool: """Python mints an Anthropic federation token when no static credential is configured.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo diff --git a/litellm/rust_bridge/responses/detectors.py b/litellm/rust_bridge/responses/detectors.py index 392c11f48a8..763b63d6742 100644 --- a/litellm/rust_bridge/responses/detectors.py +++ b/litellm/rust_bridge/responses/detectors.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # tool lists stay from pydantic import TypeAdapter import litellm +from litellm.llms.opencode.common_utils import is_opencode_api_base from litellm.rust_bridge.detectors import api_base, call_callbacks, overridden, present from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str @@ -25,7 +26,7 @@ def client_headers(request: NativeCall) -> bool: def provider_affinity(request: NativeCall) -> bool: """Python validates and fills the affinity header from the session ID.""" - return present(request, "provider_affinity_header") + return present(request, "provider_affinity_header") or is_opencode_api_base(api_base(request)) def text_format(request: NativeCall) -> bool: diff --git a/litellm/sandbox/__init__.py b/litellm/sandbox/__init__.py index 1974859fb1e..bcbb7bc5743 100644 --- a/litellm/sandbox/__init__.py +++ b/litellm/sandbox/__init__.py @@ -7,7 +7,10 @@ They resolve lazily so `import litellm` does not pull in litellm.harness. import importlib from typing import Final +from pydantic import TypeAdapter + _HARNESS_SANDBOX_MODULE: Final = "litellm.harness.sandbox" +_HARNESS_EXPORT: Final = TypeAdapter(object) _HARNESS_EXPORTS: Final = frozenset( { "local", @@ -23,5 +26,5 @@ _HARNESS_EXPORTS: Final = frozenset( def __getattr__(name: str) -> object: if name in _HARNESS_EXPORTS: - return getattr(importlib.import_module(_HARNESS_SANDBOX_MODULE), name) + return _HARNESS_EXPORT.validate_python(getattr(importlib.import_module(_HARNESS_SANDBOX_MODULE), name)) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/secret_managers/base_secret_manager.py b/litellm/secret_managers/base_secret_manager.py index e537dd1d998..9c25086c451 100644 --- a/litellm/secret_managers/base_secret_manager.py +++ b/litellm/secret_managers/base_secret_manager.py @@ -1,6 +1,6 @@ import re from abc import ABC, abstractmethod -from typing import Any, Final +from typing import Final import httpx @@ -73,7 +73,7 @@ class BaseSecretManager(ABC): optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None, tags: dict | list | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Asynchronously write a secret to the secret manager. diff --git a/litellm/secret_managers/cyberark_secret_manager.py b/litellm/secret_managers/cyberark_secret_manager.py index a1ec7c0a386..abfa9649e91 100644 --- a/litellm/secret_managers/cyberark_secret_manager.py +++ b/litellm/secret_managers/cyberark_secret_manager.py @@ -1,7 +1,7 @@ import asyncio import base64 import os -from typing import Any, Final +from typing import Final from urllib.parse import quote import httpx @@ -278,7 +278,7 @@ class CyberArkSecretManager(BaseSecretManager): optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None, tags: dict | list | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Writes a secret to CyberArk Conjur using an async HTTPX client. diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index f4234e26c60..b33cf0e6b7e 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -24,6 +24,7 @@ from litellm.types.llms.base import LiteLLMBaseModel oidc_cache: Final = DualCache() _PARSED_LITERAL: Final = TypeAdapter(object) +_SECRET_MANAGER_CLIENT: Final = TypeAdapter(object) @lru_cache(maxsize=128) @@ -383,7 +384,7 @@ def get_secret( try: if _should_read_secret_from_secret_manager() and litellm.secret_manager_client is not None: try: - client: Final = litellm.secret_manager_client + client: Final = _SECRET_MANAGER_CLIENT.validate_python(litellm.secret_manager_client) key_manager = "local" if key_management_system is not None: key_manager = key_management_system.value diff --git a/litellm/text_completions/__init__.py b/litellm/text_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py new file mode 100644 index 00000000000..ec1b96fd432 --- /dev/null +++ b/litellm/text_completions/streaming_iterator.py @@ -0,0 +1,88 @@ +from typing import Final + +from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse + + +class TextCompletionStreamWrapper: + def __init__( + self, + completion_stream, + model, + stream_options: dict | None = None, + custom_llm_provider: str | None = None, + ): + self.completion_stream = completion_stream + self.model = model + self.stream_options = stream_options + self.custom_llm_provider = custom_llm_provider + + def __iter__(self): + return self + + def __aiter__(self): + return self + + def convert_to_text_completion_object(self, chunk: ModelResponse): + try: + response: Final = TextCompletionResponse() + response["id"] = chunk.get("id", None) + response["object"] = "text_completion" + response["created"] = chunk.get("created", None) + response["model"] = chunk.get("model", None) + text_choices: Final = TextChoices() + if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices + raise Exception + + if chunk["choices"]: + delta: Final = chunk["choices"][0]["delta"] + text_choices["text"] = delta["content"] + text_choices["reasoning_content"] = delta.get("reasoning_content") + text_choices["index"] = chunk["choices"][0]["index"] + text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + response["choices"] = [text_choices] + else: + response["choices"] = [] + + # only pass usage when stream_options["include_usage"] is True + if self.stream_options and self.stream_options.get("include_usage", False) is True: + response.usage = chunk.get("usage", None) + else: + response.usage = None + + return response + except Exception as e: + raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {e}") + + def __next__(self): + # model_response = ModelResponse(stream=True, model=self.model) + TextCompletionResponse() + try: + for chunk in self.completion_stream: + if chunk == "None" or chunk is None: + raise Exception + processed_chunk = self.convert_to_text_completion_object(chunk=chunk) + return processed_chunk + raise StopIteration + except StopIteration: + raise StopIteration + except Exception as e: + from litellm.litellm_core_utils.exception_mapping_utils import exception_type + + raise exception_type( + model=self.model, + custom_llm_provider=self.custom_llm_provider or "", + original_exception=e, + completion_kwargs={}, + extra_kwargs={}, + ) + + async def __anext__(self): + try: + async for chunk in self.completion_stream: + if chunk == "None" or chunk is None: + raise Exception + processed_chunk = self.convert_to_text_completion_object(chunk=chunk) + return processed_chunk + raise StopIteration + except StopIteration: + raise StopAsyncIteration diff --git a/litellm/types/caching.py b/litellm/types/caching.py index b9bcbbb93dd..ced2065a674 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -75,7 +75,7 @@ class RedisPipelineSetOperation(TypedDict): """ key: str - value: Any + value: object ttl: int | None @@ -133,7 +133,7 @@ class HealthCheckCacheParams(LiteLLMBaseModel): host: str | None = None port: str | int | None = None - redis_kwargs: dict[str, Any] | None = None + redis_kwargs: dict[str, object] | None = None namespace: str | None = None redis_version: str | int | float | None = None diff --git a/litellm/types/completion.py b/litellm/types/completion.py index a6693030e89..0f4bdbd75ce 100644 --- a/litellm/types/completion.py +++ b/litellm/types/completion.py @@ -246,7 +246,7 @@ _CompletionDispatchContext = CompletionDispatchContext _CompletionDispatchResult = Union[ - Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]], + Coroutine[object, object, Union["ModelResponse", "CustomStreamWrapper"]], "ModelResponse", "CustomStreamWrapper", ] diff --git a/litellm/types/decisions.py b/litellm/types/decisions.py index a1d48f730c0..33b2c677d31 100644 --- a/litellm/types/decisions.py +++ b/litellm/types/decisions.py @@ -73,7 +73,7 @@ DecisionQuestion: TypeAlias = Annotated[ DecisionQuestionMap: TypeAlias = Annotated[ Mapping[Annotated[str, Field(min_length=1)], DecisionQuestion], - Field(min_length=1, max_length=MAX_DECISION_QUESTIONS), + Field(min_length=1), ] @@ -650,7 +650,7 @@ class DecisionsIRUsage: class DecisionsIRResponse: model: str | None answers: tuple[DecisionsIRAnswer, ...] - usage: DecisionsIRUsage + usage: DecisionsIRUsage | None extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) diff --git a/litellm/types/extraction.py b/litellm/types/extraction.py index 4311e549b9a..562bc41f9de 100644 --- a/litellm/types/extraction.py +++ b/litellm/types/extraction.py @@ -13,6 +13,7 @@ from pydantic import ( model_validator, with_config, ) +from pydantic.fields import FieldInfo from typing_extensions import ReadOnly, Required, Self, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -30,7 +31,7 @@ def _is_none(value: object) -> bool: return value is None -_OMITTED_WHEN_NONE: Final = Field(exclude_if=_is_none) +_OMITTED_WHEN_NONE: Final = FieldInfo.from_field(exclude_if=_is_none) OptionalJson: TypeAlias = Annotated[JsonValue, _OMITTED_WHEN_NONE] OptionalFloat: TypeAlias = Annotated[float | None, _OMITTED_WHEN_NONE] OptionalInt: TypeAlias = Annotated[int | None, _OMITTED_WHEN_NONE] diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py index 0241c9fa809..dcde0a031e3 100644 --- a/litellm/types/google_genai/main.py +++ b/litellm/types/google_genai/main.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: ToolConfigDict = _genai_types.ToolConfigDict class GenerateContentRequestDict(GenerateContentRequestParametersDict): - generationConfig: Any | None + generationConfig: object | None tools: ToolConfigDict | None class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ff063d2e233..e1b891caaaf 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -59,6 +59,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( QualifireGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.reco import ( + RecoConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import ( RepelloAIGuardrailConfigModel, ) @@ -153,6 +156,7 @@ class SupportedGuardrailIntegrations(Enum): TYPESAFE = "typesafe" STRAIKER = "straiker" ALICE = "alice" + RECO = "reco" AGENT_365 = "agent_365" LLM_SHIELD_PROXY = "llm_shield_proxy" CONDUCT = "conduct" @@ -1398,6 +1402,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o QostodianNexusConfigModel, VigilGuardGuardrailConfigModel, SingulrGuardrailConfigModel, + RecoConfigModel, Agent365GuardrailConfigModel, WingbackGuardrailConfigModel, ): diff --git a/litellm/types/images/main.py b/litellm/types/images/main.py index 5d80135a8a1..963c980a4b8 100644 --- a/litellm/types/images/main.py +++ b/litellm/types/images/main.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from typing_extensions import TypedDict @@ -20,7 +20,7 @@ class ImageEditOptionalRequestParams(TypedDict, total=False): response_format: Literal["url", "b64_json"] | None size: str | None user: str | None - imageConfig: dict[str, Any] | None + imageConfig: dict[str, object] | None class ImageEditRequestParams(ImageEditOptionalRequestParams, total=False): diff --git a/litellm/types/integrations/argilla.py b/litellm/types/integrations/argilla.py index 2c98486c7a7..754916c2f3a 100644 --- a/litellm/types/integrations/argilla.py +++ b/litellm/types/integrations/argilla.py @@ -1,10 +1,10 @@ -from typing import Any, Final +from typing import Final from typing_extensions import TypedDict class ArgillaItem(TypedDict): - fields: dict[str, Any] + fields: dict[str, object] class ArgillaPayload(TypedDict): diff --git a/litellm/types/integrations/cloudzero.py b/litellm/types/integrations/cloudzero.py index 988d501600f..dcc5d03ccca 100644 --- a/litellm/types/integrations/cloudzero.py +++ b/litellm/types/integrations/cloudzero.py @@ -1,7 +1,4 @@ -from typing import Any - - -class CBFRecord(dict[str, Any]): +class CBFRecord(dict[str, object]): """CloudZero Billing Format (CBF) record structure. This class represents a CBF record that is created from LiteLLM usage data @@ -31,4 +28,4 @@ class CBFRecord(dict[str, Any]): # Type alias for better readability in function signatures -CBFRecordDict = dict[str, Any] +CBFRecordDict = dict[str, object] diff --git a/litellm/types/integrations/compression_interception.py b/litellm/types/integrations/compression_interception.py index a1e66e5057e..e9c40fbf838 100644 --- a/litellm/types/integrations/compression_interception.py +++ b/litellm/types/integrations/compression_interception.py @@ -2,7 +2,7 @@ Type definitions for Compression Interception integration. """ -from typing import Any, Literal, TypedDict +from typing import Literal, TypedDict class CompressionInterceptionConfig(TypedDict, total=False): @@ -24,7 +24,7 @@ class CompressionInterceptionConfig(TypedDict, total=False): compression_trigger: int compression_target: int | None embedding_model: str | None - embedding_model_params: dict[str, Any] | None + embedding_model_params: dict[str, object] | None class CompressionSavingsMetadata(TypedDict): diff --git a/litellm/types/integrations/gcs_bucket.py b/litellm/types/integrations/gcs_bucket.py index 3840d7681a5..79563a0a6e4 100644 --- a/litellm/types/integrations/gcs_bucket.py +++ b/litellm/types/integrations/gcs_bucket.py @@ -31,5 +31,5 @@ class GCSLogQueueItem(TypedDict): """ payload: StandardLoggingPayload - kwargs: dict[str, Any] - response_obj: Any | None + kwargs: dict[str, object] + response_obj: object | None diff --git a/litellm/types/integrations/langsmith.py b/litellm/types/integrations/langsmith.py index 594190ddb5c..f053a5864c1 100644 --- a/litellm/types/integrations/langsmith.py +++ b/litellm/types/integrations/langsmith.py @@ -1,6 +1,6 @@ from dataclasses import dataclass from datetime import datetime -from typing import Any, NamedTuple +from typing import NamedTuple from typing_extensions import TypedDict @@ -9,7 +9,7 @@ from litellm.types.llms.base import LiteLLMBaseModel class LangsmithInputs(LiteLLMBaseModel): model: str | None = None - messages: list[Any] | None = None + messages: list[object] | None = None stream: bool | None = None call_type: str | None = None litellm_call_id: str | None = None @@ -17,9 +17,9 @@ class LangsmithInputs(LiteLLMBaseModel): temperature: float | None = None max_tokens: int | None = None custom_llm_provider: str | None = None - input: list[Any] | None = None + input: list[object] | None = None log_event_type: str | None = None - original_response: Any | None = None + original_response: object | None = None response_cost: float | None = None # LiteLLM Virtual Key specific fields diff --git a/litellm/types/integrations/posthog.py b/litellm/types/integrations/posthog.py index ac04d08107d..d9739df7df7 100644 --- a/litellm/types/integrations/posthog.py +++ b/litellm/types/integrations/posthog.py @@ -1,4 +1,4 @@ -from typing import Any, Final, TypedDict +from typing import Final, TypedDict POSTHOG_MAX_BATCH_SIZE: Final = 100 @@ -7,7 +7,7 @@ class PostHogEventPayload(TypedDict): """PostHog event payload structure""" event: str # "$ai_generation" or "$ai_embedding" - properties: dict[str, Any] + properties: dict[str, object] distinct_id: str diff --git a/litellm/types/interactions/README.md b/litellm/types/interactions/README.md index d890991d99d..45783d604db 100644 --- a/litellm/types/interactions/README.md +++ b/litellm/types/interactions/README.md @@ -9,16 +9,38 @@ https://ai.google.dev/static/api/interactions.openapi.json ### How to Regenerate -When the API spec changes, regenerate the types with: +The compliance tests in `tests/unit/interactions/test_openapi_compliance.py` read the snapshot at +`tests/unit/interactions/interactions.openapi.json`, so a PR never depends on Google's live spec. The scheduled +`interactions-spec-canary` workflow runs those tests against the live spec once a day (`--interactions-spec-url`), +files the issue "Interactions spec canary is red" on a red run and closes it once a run is green again. A red run +means the spec moved in a way these types no longer cover + +When the red run names one new status, the small fix is a hand edit adding that member to the enum the failing +schema's model carries, the way `queued` was added in #34135. When the spec moved further than that, refresh the +snapshot and regenerate the types from it: ```bash -uv tool run --from datamodel-code-generator datamodel-codegen \ - --url "https://ai.google.dev/static/api/interactions.openapi.json" \ +curl -fsSL https://ai.google.dev/static/api/interactions.openapi.json \ + -o tests/unit/interactions/interactions.openapi.json +uv tool run --from datamodel-code-generator==0.66.0 datamodel-codegen \ + --input tests/unit/interactions/interactions.openapi.json \ + --input-file-type openapi \ --output litellm/types/interactions/generated.py \ --output-model-type pydantic_v2.BaseModel \ - --target-python-version 3.9 + --target-python-version 3.10 \ + --base-class litellm.types.llms.base.LiteLLMBaseModel ``` +The pin is the same `datamodel-code-generator` version `scripts/generate_trace_types.py` uses. The generator names +inline enums by order of appearance (`Status`, `Status1`, `Status2`, ...), so a regeneration can rename them. The +compliance tests resolve each status enum through its model's `status` field, but +`tests/unit/interactions/test_background_cost_polling.py` and +`tests/unit/proxy/hooks/test_proxy_track_cost_callback.py` import `Status1` by name and need the new name, and +`litellm/types/interactions/__init__.py` exports the interaction status enum as `InteractionStatus`, which has to +point at the enum the regenerated `Interaction.status` carries (a compliance test checks their members match). The +spec has grown a lot since the 2025-12-16 generation (agents, triggers, credentials, environments), so a +regeneration today is a review-sized change rather than a refresh + Then add the LiteLLM-specific types at the bottom of the generated file: - `InteractionsAPIResponse` - `InteractionsAPIStreamingResponse` diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 5f318cc3c69..717b74e5741 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -54,6 +54,7 @@ RoutingStrategyName: TypeAlias = Literal[ TRUSTED_CALLBACK_VARS_FIELD: Final = "litellm_trusted_callback_vars" ADDRESSED_RESPONSE_ID_FIELD: Final = "_litellm_addressed_response_id" +STREAM_USAGE_INJECTED_FIELD: Final = "_litellm_stream_usage_injected" WIRE_NAME: Final = "wire_name" @@ -90,6 +91,7 @@ class ProviderConnection: litellm_credential_name: str | None = None configurable_clientside_auth_params: "Sequence[str | ConfigurableClientsideParamsCustomAuth] | None" = None use_xai_oauth: bool | None = None + chatgpt_auth_file: str | None = None github_copilot_auth_type: str | None = None github_copilot_host: str | None = None github_copilot_user_session: object | None = None @@ -385,6 +387,9 @@ class AgenticLoopState: @dataclass(frozen=True, slots=True, kw_only=True) class RouterState: weights: "RouterWeights | None" = field(default=None, metadata=wire("_router_weights")) + encrypted_content_affinity_pinned: bool | None = field( + default=None, metadata=wire("_encrypted_content_affinity_pinned") + ) fallback_depth: int | None = None max_fallbacks: int | None = None attempted_targets: "AttemptedFallbackTargets | None" = None @@ -397,6 +402,7 @@ class ProxyRequestState: trusted_callback_vars: Mapping[str, str] | None = field(default=None, metadata=wire(TRUSTED_CALLBACK_VARS_FIELD)) addressed_response_id: str | None = field(default=None, metadata=wire(ADDRESSED_RESPONSE_ID_FIELD)) strip_stream_usage: bool | None = field(default=None, metadata=wire("_litellm_strip_stream_usage")) + stream_usage_injected: bool | None = field(default=None, metadata=wire(STREAM_USAGE_INJECTED_FIELD)) client_side_timeout: bool | None = None model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 6aeead7679f..ba422a0a901 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -481,6 +481,11 @@ class ContentThinkingSignatureBlockDelta(TypedDict): signature: str +class ContentCompactionBlockDelta(TypedDict): + type: ReadOnly[Literal["compaction_delta"]] + content: ReadOnly[str | None] + + StreamingContentBlockDeltaType = Literal["text_delta", "input_json_delta", "thinking_delta", "signature_delta"] @@ -493,6 +498,7 @@ class ContentBlockDelta(TypedDict): | ContentCitationsBlockDelta | ContentThinkingBlockDelta | ContentThinkingSignatureBlockDelta + | ContentCompactionBlockDelta ) diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 03956ae7a54..9ad13243451 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -1,5 +1,5 @@ from collections.abc import Sequence -from typing import Any, Literal, TypeAlias +from typing import Literal, TypeAlias from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -20,7 +20,7 @@ class AnthropicResponseTextBlock(TypedDict, total=False): Anthropic Response Text Block: https://docs.anthropic.com/en/api/messages """ - citations: list[dict[str, Any]] | None + citations: list[dict[str, object]] | None text: str type: Literal["text"] diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index 47fec914f9c..5a2da071cec 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from openai._models import BaseModel as OpenAIObject from pydantic import ConfigDict @@ -45,7 +45,7 @@ class BaseLiteLLMOpenAIResponseObject(LiteLLMBaseModel): class HiddenParams(OpenAIObject): - original_response: str | Any | None = None + original_response: str | object | None = None model_id: str | None = None # used in Router for individual deployments api_base: str | None = None # returns api base used for making completion call _response_ms: float | None = None diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index c7023e35777..bc229421d91 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from pydantic import BaseModel from typing_extensions import ( @@ -23,7 +23,7 @@ class GenericStreamingChunk(TypedDict, total=False): class DatabricksTextContent(TypedDict, total=False): type: Literal["text"] text: Required[str] - citations: list[dict[str, Any]] | None + citations: list[dict[str, object]] | None class DatabricksReasoningSummary(TypedDict): @@ -35,7 +35,7 @@ class DatabricksReasoningSummary(TypedDict): class DatabricksReasoningContent(TypedDict, total=False): type: Literal["reasoning"] summary: Required[list[DatabricksReasoningSummary]] - citations: list[dict[str, Any]] | None + citations: list[dict[str, object]] | None AllDatabricksContentListValues = DatabricksTextContent | DatabricksReasoningContent diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 8dbc0d97d54..9a4ecf37127 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -55,6 +55,7 @@ except (ImportError, AttributeError): from typing import Annotated +from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ( Reasoning, ResponseIncludable, @@ -682,12 +683,12 @@ class ChatCompletionCachedContent(TypedDict): class PromptCacheBreakpoint(TypedDict): mode: ReadOnly[Literal["explicit"]] - ttl: NotRequired[ReadOnly[Literal["30m"]]] + ttl: NotRequired[ReadOnly[Literal["30m", "12h"]]] class PromptCacheOptions(TypedDict, total=False): mode: ReadOnly[Literal["implicit", "explicit"]] - ttl: ReadOnly[Literal["30m"]] + ttl: ReadOnly[Literal["30m", "12h"]] class ChatCompletionThinkingBlock(TypedDict, total=False): @@ -1381,10 +1382,23 @@ class OutputTokensDetails(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} +class CacheWriteTokensDetails(BaseLiteLLMOpenAIResponseObject): + ttl_30m_tokens: int + ttl_12h_tokens: int + + model_config = ConfigDict(extra="allow", frozen=True) + + +def _is_none(value: object) -> bool: + return value is None + + class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 cached_tokens_details: CachedTokensDetails | None = None + cache_write_tokens: int | None = Field(default=None, exclude_if=_is_none) + cache_write_tokens_details: CacheWriteTokensDetails | None = Field(default=None, exclude_if=_is_none) image_tokens: int | None = None text_tokens: int | None = None video_tokens: int | None = None @@ -1488,6 +1502,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ] ) parallel_tool_calls: bool | None = None diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 0480d6b1f3e..f1176f777c7 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -239,8 +239,18 @@ class VoiceConfig(TypedDict): prebuiltVoiceConfig: PrebuiltVoiceConfig +class SpeakerVoiceConfig(TypedDict): + speaker: ReadOnly[str] + voiceConfig: ReadOnly[VoiceConfig] + + +class MultiSpeakerVoiceConfig(TypedDict): + speakerVoiceConfigs: ReadOnly[Sequence[SpeakerVoiceConfig]] + + class SpeechConfig(TypedDict, total=False): voiceConfig: VoiceConfig + multiSpeakerVoiceConfig: ReadOnly[MultiSpeakerVoiceConfig] languageCode: str diff --git a/litellm/types/llms/vertex_ai_text_to_speech.py b/litellm/types/llms/vertex_ai_text_to_speech.py index e8f9fbe4a24..86b131dce59 100644 --- a/litellm/types/llms/vertex_ai_text_to_speech.py +++ b/litellm/types/llms/vertex_ai_text_to_speech.py @@ -4,7 +4,9 @@ Type definitions for Vertex AI Text-to-Speech API Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/synthesize """ -from typing_extensions import TypedDict +from collections.abc import Sequence + +from typing_extensions import ReadOnly, TypedDict class VertexTextToSpeechInput(TypedDict, total=False): @@ -16,6 +18,16 @@ class VertexTextToSpeechInput(TypedDict, total=False): text: str | None ssml: str | None + prompt: ReadOnly[str | None] + + +class VertexTextToSpeechSpeakerVoiceConfig(TypedDict): + speakerAlias: ReadOnly[str] + speakerId: ReadOnly[str] + + +class VertexTextToSpeechMultiSpeakerVoiceConfig(TypedDict): + speakerVoiceConfigs: ReadOnly[Sequence[VertexTextToSpeechSpeakerVoiceConfig]] class VertexTextToSpeechVoice(TypedDict, total=False): @@ -29,6 +41,8 @@ class VertexTextToSpeechVoice(TypedDict, total=False): languageCode: str name: str + modelName: ReadOnly[str] + multiSpeakerVoiceConfig: ReadOnly[VertexTextToSpeechMultiSpeakerVoiceConfig] class VertexTextToSpeechAudioConfig(TypedDict, total=False): @@ -42,6 +56,7 @@ class VertexTextToSpeechAudioConfig(TypedDict, total=False): audioEncoding: str speakingRate: str + sampleRateHertz: ReadOnly[int] class VertexTextToSpeechRequest(TypedDict, total=False): diff --git a/litellm/types/management_endpoints/cache_settings_endpoints.py b/litellm/types/management_endpoints/cache_settings_endpoints.py index f6394f52c99..3584f8ce1da 100644 --- a/litellm/types/management_endpoints/cache_settings_endpoints.py +++ b/litellm/types/management_endpoints/cache_settings_endpoints.py @@ -2,7 +2,7 @@ Types and field definitions for cache settings management endpoints """ -from typing import Any, Final +from typing import Final from litellm.types.llms.base import LiteLLMBaseModel @@ -10,9 +10,9 @@ from litellm.types.llms.base import LiteLLMBaseModel class CacheSettingsField(LiteLLMBaseModel): field_name: str field_type: str - field_value: Any + field_value: object field_description: str - field_default: Any = None + field_default: object = None options: list[str] | None = None # For fields with predefined options/enum values ui_field_name: str # User-friendly display name link: str | None = None # Documentation link for the field diff --git a/litellm/types/management_endpoints/router_settings_endpoints.py b/litellm/types/management_endpoints/router_settings_endpoints.py index 9c8fe446d72..64068f30a16 100644 --- a/litellm/types/management_endpoints/router_settings_endpoints.py +++ b/litellm/types/management_endpoints/router_settings_endpoints.py @@ -2,7 +2,7 @@ Types and field definitions for router settings management endpoints """ -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import Field, field_validator @@ -72,9 +72,9 @@ class FallbackDeleteResponse(LiteLLMBaseModel): class RouterSettingsField(LiteLLMBaseModel): field_name: str field_type: str - field_value: Any + field_value: object field_description: str - field_default: Any = None + field_default: object = None options: list[str] | None = None # For fields with predefined options/enum values ui_field_name: str # User-friendly display name link: str | None = None # Documentation link for the field diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index f0be66f055d..3fa9adef913 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -113,7 +113,7 @@ class MCPPublicServer(LiteLLMBaseModel): transport: MCPTransportType spec_path: str | None = None auth_type: MCPAuthType | None = None - mcp_info: dict[str, Any] | None = None + mcp_info: dict[str, object] | None = None class MCPAllowedClient(LiteLLMBaseModel): @@ -448,11 +448,11 @@ class MCPPreCallRequestObject(LiteLLMBaseModel): """ tool_name: str - arguments: dict[str, Any] + arguments: dict[str, object] server_name: str | None = None tool_description: str | None = None tool_input_schema: Mapping[str, object] | None = None - user_api_key_auth: dict[str, Any] | None = None + user_api_key_auth: dict[str, object] | None = None hidden_params: HiddenParams = HiddenParams() @@ -462,7 +462,7 @@ class MCPPreCallResponseObject(LiteLLMBaseModel): """ should_proceed: bool = True - modified_arguments: dict[str, Any] | None = None + modified_arguments: dict[str, object] | None = None error_message: str | None = None hidden_params: HiddenParams = HiddenParams() @@ -473,7 +473,7 @@ class MCPDuringCallRequestObject(LiteLLMBaseModel): """ tool_name: str - arguments: dict[str, Any] + arguments: dict[str, object] server_name: str | None = None tool_description: str | None = None tool_input_schema: Mapping[str, object] | None = None diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 219f23200d2..d29e4db55a4 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -117,7 +117,7 @@ class MCPServer(LiteLLMBaseModel): # Admin-configured env vars. Each entry is {name, value, scope, description}. # scope=="global" values are interpolated into static_headers using ${NAME}. # scope=="user" values must be supplied per-user. - env_vars: list[dict[str, Any]] | None = None + env_vars: list[dict[str, object]] | None = None # OAuth-specific fields authorization_server_metadata: Mapping[str, object] | None = Field(default=None, exclude=True) client_id: str | None = None @@ -240,7 +240,7 @@ class MCPServer(LiteLLMBaseModel): # token_validation: key-value pairs that must match fields in the OAuth token # response (supports dot-notation for nested fields, e.g. "team.enterprise_id"). # Tokens that fail validation are rejected before storage. - token_validation: dict[str, Any] | None = None + token_validation: dict[str, object] | None = None oauth_identity_binding: MCPOAuthIdentityBinding | None = None # Optional TTL override (seconds) for the Redis per-user token cache, capped # at the token's expires_in minus the expiry buffer so a cached entry never diff --git a/litellm/types/prompts/init_prompts.py b/litellm/types/prompts/init_prompts.py index fc4c36e760a..c1a3ab56311 100644 --- a/litellm/types/prompts/init_prompts.py +++ b/litellm/types/prompts/init_prompts.py @@ -1,6 +1,6 @@ from datetime import datetime from enum import Enum -from typing import Any, Literal +from typing import Literal from pydantic import ConfigDict @@ -31,7 +31,7 @@ class PromptLiteLLMParams(LiteLLMBaseModel): api_base: str | None = None api_key: str | None = None - provider_specific_query_params: dict[str, Any] | None = None + provider_specific_query_params: dict[str, object] | None = None ignore_prompt_manager_model: bool | None = False ignore_prompt_manager_optional_params: bool | None = False @@ -66,7 +66,7 @@ class PromptSpec(LiteLLMBaseModel): class PromptTemplateBase(LiteLLMBaseModel): litellm_prompt_id: str content: str - metadata: dict[str, Any] | None = None + metadata: dict[str, object] | None = None class PromptInfoResponse(LiteLLMBaseModel): diff --git a/litellm/types/proxy/callback_logs_endpoints.py b/litellm/types/proxy/callback_logs_endpoints.py index b2519d97122..0291cc902bc 100644 --- a/litellm/types/proxy/callback_logs_endpoints.py +++ b/litellm/types/proxy/callback_logs_endpoints.py @@ -5,7 +5,7 @@ External producers (e.g. the litellm-rust gateway) POST finished logging payloads here; the proxy replays them through the standard callback fan-out. """ -from typing import Any, Literal +from typing import Literal from pydantic import Field @@ -17,7 +17,7 @@ class CallbackLogRecord(LiteLLMBaseModel): """A single finished logging event to replay through the callbacks.""" status: Literal["success", "failure"] - standard_logging_payload: dict[str, Any] + standard_logging_payload: dict[str, object] error: str | None = None diff --git a/litellm/types/proxy/cloudzero_endpoints.py b/litellm/types/proxy/cloudzero_endpoints.py index ece294674b9..f87537494bd 100644 --- a/litellm/types/proxy/cloudzero_endpoints.py +++ b/litellm/types/proxy/cloudzero_endpoints.py @@ -46,7 +46,7 @@ class CloudZeroExportResponse(LiteLLMBaseModel): dry_run_data: dict[str, Any] | None = Field( None, description="Dry run data including usage data and CBF transformed data" ) - summary: dict[str, Any] | None = Field(None, description="Summary statistics for dry run") + summary: dict[str, object] | None = Field(None, description="Summary statistics for dry run") class CloudZeroSettingsView(LiteLLMBaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py index 60846b2a1bd..4f4bebf9704 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py @@ -1,5 +1,3 @@ -from typing import Any - from pydantic import Field from typing_extensions import TypedDict @@ -23,7 +21,7 @@ class AzurePromptShieldGuardrailResponse(TypedDict): """Configuration parameters for the Azure Prompt Shield guardrail""" userPromptAnalysis: UserPromptAnalysis - documentsAnalysis: list[dict[str, Any]] + documentsAnalysis: list[dict[str, object]] class AzurePromptShieldGuardrailConfigModel( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py index 90662e84ed5..6ee5741baf5 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py @@ -1,4 +1,4 @@ -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import Field from typing_extensions import Required, TypedDict @@ -37,7 +37,7 @@ class AzureTextModerationGuardrailResponseCategoriesAnalysis(TypedDict): class AzureTextModerationGuardrailResponse(TypedDict): """Response from the Azure Text Moderation guardrail""" - blocklistsMatch: list[dict[str, Any]] + blocklistsMatch: list[dict[str, object]] categoriesAnalysis: list[AzureTextModerationGuardrailResponseCategoriesAnalysis] diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py b/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py index 6598f1de62b..c0f5ea5346e 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from pydantic import Field @@ -93,7 +93,7 @@ class CompresrGuardrailOptionalParams(LiteLLMBaseModel): "Unset lets the server default apply (~10.0)." ), ) - compression_params: dict[str, Any] | None = Field( + compression_params: dict[str, object] | None = Field( default=None, description=( "Passthrough of extra parameters forwarded verbatim in the Compresr " diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py b/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py index 78800a5e9a2..d49a3afae04 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Any, Literal, TypedDict +from typing import Literal, TypedDict from pydantic import Field @@ -77,7 +77,7 @@ class CompetitorIntentDetection(TypedDict): action_hint: str entities: dict[str, list[str]] signals: list[str] - evidence: list[dict[str, Any]] + evidence: list[dict[str, object]] ContentFilterDetection = PatternDetection | BlockedWordDetection | CategoryKeywordDetection | CompetitorIntentDetection @@ -157,7 +157,7 @@ class LitellmContentFilterGuardrailConfigModel(GuardrailConfigModel): ) # Competitor intent blocker (generic; industry presets add domain_words, etc.) - competitor_intent_config: dict[str, Any] | None = Field( + competitor_intent_config: dict[str, object] | None = Field( default=None, description="Optional config for intent-based competitor comparison detection. " "Keys: brand_self (list), competitors (list), competitor_aliases (dict), " diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/types/proxy/guardrails/guardrail_hooks/presidio.py index 9d1e375af72..6fe324ce6b0 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/presidio.py @@ -1,6 +1,6 @@ -from typing import Any +from collections.abc import Mapping -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.guardrails import PiiEntityType @@ -17,5 +17,5 @@ class PresidioAnalyzeResponseItem(TypedDict, total=False): start: int | None end: int | None score: float | None - analysis_explanation: dict[str, Any] | None - recognition_metadata: dict[str, Any] | None + analysis_explanation: ReadOnly[Mapping[str, object] | None] + recognition_metadata: ReadOnly[Mapping[str, object] | None] diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/reco.py b/litellm/types/proxy/guardrails/guardrail_hooks/reco.py new file mode 100644 index 00000000000..ef3b3935340 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/reco.py @@ -0,0 +1,48 @@ +import uuid +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from .base import GuardrailConfigModel + + +def validate_reco_tenant_id(value: str) -> str: + try: + uuid.UUID(value) + except ValueError as e: + raise ValueError(f"reco_tenant_id must be a valid UUID, got {value!r}") from e + return value + + +class RecoOptionalParams(BaseModel): + """Configuration parameters for the Reco guardrail""" + + model_config = ConfigDict(frozen=True) + + reco_tenant_id: str = Field( + description="Tenant identifier for the Reco account, as a UUID. Sent as the X-Reco-Tenant-Id header on every guardrail request.", + ) + api_base: str = Field( + description="Base URL for the Reco guardrail API. Reco's endpoint is per-region and per-silo, so this must be set explicitly rather than defaulted.", + ) + unreachable_fallback: Literal["fail_closed", "fail_open"] = Field( + default="fail_closed", + description="Behavior when Reco is unreachable. 'fail_closed' (default) blocks the request; 'fail_open' lets it through uninspected. Only applies while fail_on_error is true.", + ) + fail_on_error: bool = Field( + default=True, + description="Behavior on any Reco error, including unreachability. True (default) blocks the request; False lets it through uninspected.", + ) + + @field_validator("reco_tenant_id") + @classmethod + def _check_reco_tenant_id(cls, v: str) -> str: + return validate_reco_tenant_id(v) + + +class RecoConfigModel(GuardrailConfigModel[RecoOptionalParams]): + """Configuration parameters for the Reco guardrail""" + + @staticmethod + def ui_friendly_name() -> str: + return "Reco" diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 66702abfb34..213c911cfdc 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -1,7 +1,6 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import date from enum import Enum -from typing import Any from pydantic import ConfigDict, Field from typing_extensions import TypedDict @@ -61,7 +60,7 @@ class KeyMetricWithMetadata(MetricBase): class MetricWithMetadata(MetricBase): - metadata: dict[str, Any] = Field(default_factory=dict) + metadata: Mapping[str, object] = Field(default_factory=dict) # API key breakdown for this metric (e.g., which API keys are using this MCP server) api_key_breakdown: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} diff --git a/litellm/types/proxy/management_endpoints/config_overrides.py b/litellm/types/proxy/management_endpoints/config_overrides.py index c82e616cf9a..b807b5537d5 100644 --- a/litellm/types/proxy/management_endpoints/config_overrides.py +++ b/litellm/types/proxy/management_endpoints/config_overrides.py @@ -1,5 +1,3 @@ -from typing import Any - from pydantic import Field from litellm.types.llms.base import LiteLLMBaseModel @@ -103,5 +101,5 @@ class ConfigOverrideSettingsResponse(LiteLLMBaseModel): """Response model for config override settings GET endpoints.""" config_type: str = Field(description="The type of config override") - values: dict[str, Any] = Field(description="Current configuration values (sensitive fields decrypted)") - field_schema: dict[str, Any] = Field(description="Schema information for UI rendering") + values: dict[str, object] = Field(description="Current configuration values (sensitive fields decrypted)") + field_schema: dict[str, object] = Field(description="Schema information for UI rendering") diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index febd4485d10..fc0b73f77b6 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -1,5 +1,5 @@ from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import ConfigDict, Field, field_validator from typing_extensions import ReadOnly, TypedDict @@ -80,7 +80,7 @@ class UserUpdateResult(LiteLLMBaseModel): user_email: str | None = None success: bool error: str | None = None - updated_user: dict[str, Any] | None = None + updated_user: dict[str, object] | None = None class BulkUpdateUserResponse(LiteLLMBaseModel): diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index baa02b68d98..172357d8ac4 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -55,6 +55,8 @@ class LiteLLMKeyType(str, enum.Enum): class GenerateKeyRequest(KeyRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) soft_budget: float | None = None send_invite_email: bool | None = None key_type: LiteLLMKeyType | None = Field( @@ -82,6 +84,8 @@ def _is_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: class UpdateKeyRequest(KeyRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) # Note: the defaults of all Params here MUST BE NONE # else they will get overwritten duration: str | None = None diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index 6a4160c3c50..c012f872280 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,6 +1,5 @@ from collections.abc import Mapping from datetime import datetime -from typing import Any from pydantic import ConfigDict, Field @@ -22,7 +21,7 @@ class UpdateUsefulLinksRequest(LiteLLMBaseModel): # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) - useful_links: dict[str, str | dict[str, Any]] + useful_links: dict[str, str | dict[str, object]] class AutoRouterClassifierDefaultPromptResponse(LiteLLMBaseModel): diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 743b2d7c963..742842bd028 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,5 +1,6 @@ from collections.abc import Mapping -from typing import Any, Final, Literal +from datetime import datetime +from typing import Final, Literal from pydantic import ConfigDict, Field, field_validator, model_validator @@ -83,11 +84,22 @@ class BulkUpdateTeamMemberPermissionsResponse(LiteLLMBaseModel): permissions_appended: list[str] | None = None +class TeamListCallerMembership(LiteLLMBaseModel): + """The calling user's own spend and effective member budget in a team.""" + + spend: float = 0.0 + max_budget: float | None = None + budget_reset_at: datetime | None = None + + model_config = ConfigDict(frozen=True) + + class TeamListItem(LiteLLM_TeamTable): """A team item in the paginated list response, enriched with computed fields.""" members_count: int = 0 keys_count: int = 0 + caller_membership: TeamListCallerMembership | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None @@ -121,8 +133,8 @@ class TeamMemberAddResult(LiteLLMBaseModel): user_email: str | None = None success: bool error: str | None = None - updated_user: dict[str, Any] | None = None - updated_team_membership: dict[str, Any] | None = None + updated_user: dict[str, object] | None = None + updated_team_membership: dict[str, object] | None = None class BulkTeamMemberAddResponse(LiteLLMBaseModel): @@ -133,7 +145,7 @@ class BulkTeamMemberAddResponse(LiteLLMBaseModel): total_requested: int successful_additions: int failed_additions: int - updated_team: dict[str, Any] | None = None + updated_team: dict[str, object] | None = None class TeamMemberRef(MemberDeleteRequest): diff --git a/litellm/types/proxy/management_endpoints/user_requests.py b/litellm/types/proxy/management_endpoints/user_requests.py index 49dda5995bc..e53dc4152c3 100644 --- a/litellm/types/proxy/management_endpoints/user_requests.py +++ b/litellm/types/proxy/management_endpoints/user_requests.py @@ -16,6 +16,8 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): class NewUserRequest(GenerateRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) max_budget: float | None = None user_email: str | None = None user_alias: str | None = None @@ -47,6 +49,8 @@ class NewUserRequest(GenerateRequestBase): class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) # repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model password: str | None = Field(default=None, repr=False) spend: float | None = None diff --git a/litellm/types/proxy/prompt_endpoints.py b/litellm/types/proxy/prompt_endpoints.py index 930713663a1..cded9f5fb87 100644 --- a/litellm/types/proxy/prompt_endpoints.py +++ b/litellm/types/proxy/prompt_endpoints.py @@ -1,9 +1,7 @@ -from typing import Any - from litellm.types.llms.base import LiteLLMBaseModel class TestPromptRequest(LiteLLMBaseModel): dotprompt_content: str - prompt_variables: dict[str, Any] | None = None + prompt_variables: dict[str, object] | None = None conversation_history: list[dict[str, str]] | None = None diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index c487ce8b6c5..0e8a2bb3401 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -1,5 +1,5 @@ from collections.abc import Mapping, Sequence -from typing import Any, Literal +from typing import Literal from pydantic import ConfigDict @@ -13,7 +13,7 @@ class PublicModelHubInfo(LiteLLMBaseModel): # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) - useful_links: dict[str, str | dict[str, Any]] | None + useful_links: dict[str, str | dict[str, object]] | None class ProviderCredentialField(LiteLLMBaseModel): diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py index 1833d8b8386..ccf050a8921 100644 --- a/litellm/types/proxy/vantage_endpoints.py +++ b/litellm/types/proxy/vantage_endpoints.py @@ -2,8 +2,8 @@ Vantage endpoint types for LiteLLM Proxy """ +from collections.abc import Mapping from datetime import datetime -from typing import Any from pydantic import Field, field_validator @@ -57,10 +57,10 @@ class VantageExportResponse(LiteLLMBaseModel): message: str status: str - dry_run_data: dict[str, Any] | None = Field( + dry_run_data: Mapping[str, object] | None = Field( None, description="Dry run data including usage data and FOCUS transformed data" ) - summary: dict[str, Any] | None = Field(None, description="Summary statistics for dry run") + summary: Mapping[str, object] | None = Field(None, description="Summary statistics for dry run") class VantageSettingsView(LiteLLMBaseModel): diff --git a/litellm/types/rag.py b/litellm/types/rag.py index 2f912d21587..02c06f242d1 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -3,7 +3,7 @@ Type definitions for RAG (Retrieval Augmented Generation) Ingest API. """ from collections.abc import Mapping -from typing import Any, Literal +from typing import Literal from pydantic import ConfigDict from typing_extensions import ReadOnly, TypedDict @@ -232,7 +232,7 @@ class RAGIngestRequest(LiteLLMBaseModel): file_url: str | None = None # URL to fetch file from file_id: str | None = None # Existing file ID - ingest_options: dict[str, Any] # RAGIngestOptions as dict for flexibility + ingest_options: dict[str, object] # RAGIngestOptions as dict for flexibility model_config = ConfigDict(extra="allow") # Allow additional fields @@ -260,7 +260,7 @@ class RAGQueryRequest(LiteLLMBaseModel): """Request body for RAG query API.""" model: str - messages: list[Any] + messages: list[object] retrieval_config: RAGRetrievalConfig rerank: RAGRerankConfig | None = None stream: bool | None = False diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 5588527c141..72dc63a581f 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -5,7 +5,7 @@ from typing import Final, Literal, Optional, Union from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.response_function_web_search import ActionSearchSource, ResponseFunctionWebSearch from pydantic import ConfigDict, PrivateAttr -from typing_extensions import Any, TypedDict +from typing_extensions import TypedDict from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject @@ -100,7 +100,7 @@ class OutputCodeInterpreterCall(BaseLiteLLMOpenAIResponseObject): def build_code_interpreter_log_outputs( - content: Any, + content: object, ) -> list[OutputCodeInterpreterCallLog] | None: """Convert Anthropic bash_code_execution stdout/stderr to log outputs. diff --git a/litellm/types/router.py b/litellm/types/router.py index cd1389aea9d..ba9b82e2a57 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -572,6 +572,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): default=False, description="Use stored xAI OAuth credentials when no xAI API key is configured.", ) + chatgpt_auth_file: str | None = Field( + default=None, + description="Path to a ChatGPT OAuth auth.json for this deployment.", + ) fireworks_forward_user_id: bool | None = Field( default=None, description="Send the LiteLLM user id of the calling key as the `user` field on Fireworks AI chat, responses and messages requests.", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5939a350970..4f39741a079 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -322,11 +322,23 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_above_272k_tokens_flex: float | None cache_creation_input_token_cost_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_creation_input_token_cost_above_1hr: float | None + cache_creation_input_token_cost_above_12hr: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_flex: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_priority: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_above_272k_tokens: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_batches: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_ultrafast: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_creation_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_creation_input_token_cost_priority: float | None # OpenAI priority service tier pricing cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_read_input_token_cost: float | None cache_read_input_audio_token_cost: ReadOnly[float | None] + cache_read_input_audio_token_cost_flex: ReadOnly[float | None] + cache_read_input_audio_token_cost_priority: ReadOnly[float | None] cache_read_input_image_token_cost: ReadOnly[float | None] explicit_cache_read_input_token_cost: ReadOnly[float | None] cache_read_input_token_cost_flex: float | None # OpenAI flex service tier pricing @@ -342,6 +354,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_read_input_token_cost_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_read_input_token_cost_above_512k_tokens: float | None cache_read_input_token_cost_batches: ReadOnly[float | None] + cache_read_input_audio_token_cost_batches: ReadOnly[float | None] cache_read_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] cache_read_input_token_cost_above_100k_tokens_batches: ReadOnly[float | None] cache_read_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] @@ -355,6 +368,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): off_peak_pricing: ReadOnly[OffPeakPricing | None] # time-windowed off-peak rates input_cost_per_character: float | None # only for vertex ai models input_cost_per_audio_token: float | None + input_cost_per_audio_token_flex: ReadOnly[float | None] + input_cost_per_audio_token_priority: ReadOnly[float | None] input_cost_per_token_above_128k_tokens: float | None # only for vertex ai models input_cost_per_token_above_200k_tokens: float | None # only for vertex ai gemini-2.5-pro models input_cost_per_token_above_100k_tokens: ReadOnly[float | None] @@ -1804,7 +1819,9 @@ class CompletionTokensDetailsWrapper(CompletionTokensDetails): # wrapper for ol class CacheCreationTokenDetails(LiteLLMBaseModel): ephemeral_5m_input_tokens: int | None = None + ephemeral_30m_input_tokens: int | None = None ephemeral_1h_input_tokens: int | None = None + ephemeral_12h_input_tokens: int | None = None class PromptTokensDetailsWrapper( @@ -3954,6 +3971,16 @@ class MirroredPricingParams(LiteLLMBaseModel): explicit_cache_read_input_token_cost: float | None = None cache_creation_input_token_cost: float | None = None cache_creation_input_token_cost_above_1hr: float | None = None + cache_creation_input_token_cost_above_12hr: float | None = None + cache_creation_input_token_cost_above_12hr_flex: float | None = None + cache_creation_input_token_cost_above_12hr_priority: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority: float | None = None + cache_creation_input_token_cost_above_12hr_batches: float | None = None + cache_creation_input_token_cost_above_12hr_ultrafast: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast: float | None = None tiered_pricing: list[dict[str, Any]] | None = None @@ -4006,14 +4033,19 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_read_input_token_cost_above_100k_tokens_batches: float | None = None cache_read_input_token_cost_above_200k_tokens_batches: float | None = None cache_read_input_token_cost_above_272k_tokens_batches: float | None = None + cache_read_input_audio_token_cost_batches: float | None = None cache_creation_input_token_cost_batches: float | None = None cache_creation_input_token_cost_above_100k_tokens_batches: float | None = None cache_creation_input_token_cost_above_200k_tokens_batches: float | None = None cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None cache_read_input_audio_token_cost: float | None = None + cache_read_input_audio_token_cost_flex: float | None = None + cache_read_input_audio_token_cost_priority: float | None = None cache_read_input_image_token_cost: float | None = None input_cost_per_character_above_128k_tokens: float | None = None input_cost_per_audio_token: float | None = None + input_cost_per_audio_token_flex: float | None = None + input_cost_per_audio_token_priority: float | None = None input_cost_per_token_cache_hit: float | None = None input_cost_per_token_above_128k_tokens: float | None = None input_cost_per_token_above_100k_tokens: float | None = None @@ -4328,6 +4360,7 @@ server_owned_wif_litellm_params: Final = ( + openai_wif_litellm_params + oauth_token_exchange_litellm_params + github_copilot_oauth_litellm_params + + ("chatgpt_auth_file",) ) secret_bearing_wif_litellm_params: Final = tuple(sorted(WIF_SECRET_BEARING_KEYS)) @@ -4454,6 +4487,7 @@ class LlmProviders(str, Enum): PUBLICAI = "publicai" V0 = "v0" MORPH = "morph" + OPENCODE_GO = "opencode_go" LAMBDA_AI = "lambda_ai" INCEPTION = "inception" TEXT_COMPLETION_INCEPTION = "text-completion-inception" diff --git a/litellm/types/vector_store_files.py b/litellm/types/vector_store_files.py index 6b4953965d5..3623746216a 100644 --- a/litellm/types/vector_store_files.py +++ b/litellm/types/vector_store_files.py @@ -1,7 +1,9 @@ +import builtins +from collections.abc import Mapping from enum import Enum -from typing import Any, Literal +from typing import Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class VectorStoreFileStatus(str, Enum): @@ -33,7 +35,7 @@ class VectorStoreFileObject(TypedDict, total=False): usage_bytes: int | None vector_store_id: str status: VectorStoreFileStatus - last_error: dict[str, Any] | None + last_error: ReadOnly[Mapping[str, builtins.object] | None] chunking_strategy: VectorStoreFileChunkingStrategy | None attributes: dict[str, str | int | float | bool] | None @@ -83,5 +85,5 @@ class VectorStoreFileContentResponse(TypedDict, total=False): class VectorStoreFileAuthCredentials(TypedDict, total=False): - headers: dict[str, Any] - query_params: dict[str, Any] + headers: ReadOnly[Mapping[str, object]] + query_params: ReadOnly[Mapping[str, object]] diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 5e7d517ece2..2259339e479 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -31,7 +31,7 @@ class LiteLLM_ManagedVectorStore(TypedDict, total=False): vector_store_name: str | None vector_store_description: str | None - vector_store_metadata: dict[str, Any] | str | None + vector_store_metadata: dict[str, object] | str | None created_at: datetime | None updated_at: datetime | None diff --git a/litellm/utils.py b/litellm/utils.py index 897138a940c..1cc7c63e6a5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -71,6 +71,7 @@ from litellm.constants import ( DEFAULT_TRIM_RATIO, FUNCTION_DEFINITION_TOKEN_COUNT, HF_CONFIG_FETCH_TIMEOUT_SECONDS, + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, INITIAL_RETRY_DELAY, JITTER, MAX_RETRY_DELAY, @@ -92,6 +93,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) 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.offline_mode import is_offline_mode from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens from litellm.litellm_core_utils.tool_fields import MalformedField, resolve_deployment_tool_fields @@ -217,6 +219,7 @@ def _get_cached_audio_utils(): return _audio_utils_module +from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper as _TextCompletionStreamWrapper from litellm.types.llms.openai import ( AllMessageValues, AllPromptValues, @@ -261,13 +264,17 @@ from litellm.types.utils import ( SearchProviders, SelectTokenizerResponse, StreamingChoices, - TextChoices, - TextCompletionResponse, TranscriptionResponse, Usage, credential_api_key_secret_reference, is_litellm_owned_kwarg, ) +from litellm.types.utils import ( + TextChoices as _TextChoices, +) +from litellm.types.utils import ( + TextCompletionResponse as _TextCompletionResponse, +) _CALL_TYPE_ENUM_MAP: Final[dict] = {ct.value: ct for ct in CallTypes} _BACKFILL_MODES: Final = frozenset({"chat", "responses"}) @@ -298,10 +305,14 @@ except (ImportError, AttributeError, TypeError): claude_json_str = json.dumps(json_data) import importlib.metadata from collections.abc import AsyncIterator, Callable, Collection, Iterable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable from typing_extensions import assert_never +TextChoices: TypeAlias = _TextChoices +TextCompletionResponse: TypeAlias = _TextCompletionResponse +TextCompletionStreamWrapper: TypeAlias = _TextCompletionStreamWrapper + from litellm import utils as litellm_utils # These are lazy loaded via __getattr__ @@ -339,7 +350,8 @@ if TYPE_CHECKING: # Type stubs for lazy-loaded functions to help mypy understand their types # These imports allow mypy to understand the types when these are accessed via __getattr__ - from litellm.litellm_core_utils.exception_mapping_utils import exception_type + exception_type: Callable[..., Exception] + from litellm.litellm_core_utils.get_litellm_params import ( get_base_model_from_litellm_call_metadata, get_litellm_params, @@ -859,6 +871,11 @@ def _is_gemini_model(model: str | None, custom_llm_provider: str | None) -> bool return model is not None and "gemini" in model.lower() +def is_gemini_tts_model(model: str) -> bool: + model_name: Final = model.lower() + return "gemini" in model_name and "tts" in model_name + + def _remove_thought_signature_from_id(tool_call_id: str, separator: str) -> str: """ Remove thought signature from a tool call ID. @@ -1125,83 +1142,90 @@ def function_setup( dynamic_callbacks: Final[list[str | Callable | CustomLogger] | None] = kwargs.pop("callbacks", None) all_callbacks: Final = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) - if len(all_callbacks) > 0: - for callback in all_callbacks: - # check if callback is a string - e.g. "lago", "openmeter" - if isinstance(callback, str): - callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( - callback, - internal_usage_cache=None, - llm_router=None, - ) - if callback is None or any( - type(cb) is type(callback) for cb in litellm._async_success_callback - ): # don't double add a callback - continue - if callback not in litellm.input_callback: - litellm.input_callback.append(callback) - if callback not in litellm.success_callback: - litellm.logging_callback_manager.add_litellm_success_callback(callback) - if callback not in litellm.failure_callback: - litellm.logging_callback_manager.add_litellm_failure_callback(callback) - if callback not in litellm._async_success_callback: - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - if callback not in litellm._async_failure_callback: - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - print_verbose(f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}") + with litellm.logging_callback_manager.callback_lists_lock: + if len(all_callbacks) > 0: + for callback in all_callbacks: + # check if callback is a string - e.g. "lago", "openmeter" + if isinstance(callback, str): + callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( + callback, + internal_usage_cache=None, + llm_router=None, + ) + if callback is None or any( + type(cb) is type(callback) for cb in litellm._async_success_callback + ): # don't double add a callback + continue + if callback not in litellm.input_callback: + litellm.input_callback.append(callback) + if callback not in litellm.success_callback: + litellm.logging_callback_manager.add_litellm_success_callback(callback) + if callback not in litellm.failure_callback: + litellm.logging_callback_manager.add_litellm_failure_callback(callback) + if callback not in litellm._async_success_callback: + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + if callback not in litellm._async_failure_callback: + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + print_verbose( + f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}" + ) - if ( - len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0 - ) and len(callback_list) == 0: - callback_list = list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) + callbacks_to_set_up: Final = ( + list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) + if len(callback_list) == 0 + else [] + ) + if callbacks_to_set_up: + callback_list = callbacks_to_set_up + ## ASYNC CALLBACKS - safety net for callbacks added via direct append + if len(litellm.input_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.input_callback): + if coroutine_checker.is_async_callable(callback): + litellm._async_input_callback.append(callback) + removed_async_items.append(index) + + # Pop the async items from input_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.input_callback.pop(index) + if len(litellm.success_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.success_callback): + if coroutine_checker.is_async_callable(callback): + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + removed_async_items.append(index) + elif callback == "dynamodb" or callback == "openmeter": + # dynamo is an async callback, it's used for the proxy and needs to be async + # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + removed_async_items.append(index) + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): + add_custom_logger_callback_to_specific_event(callback, "success") + + # Pop the async items from success_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.success_callback.pop(index) + + if len(litellm.failure_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.failure_callback): + if coroutine_checker.is_async_callable(callback): + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + removed_async_items.append(index) + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): + add_custom_logger_callback_to_specific_event(callback, "failure") + + # Pop the async items from failure_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.failure_callback.pop(index) + if callbacks_to_set_up: get_set_callbacks: Final[Callable[[], Callable[..., None]]] = litellm_utils.get_set_callbacks - get_set_callbacks()(callback_list=callback_list, function_id=function_id) - ## ASYNC CALLBACKS - safety net for callbacks added via direct append - if len(litellm.input_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.input_callback): - if coroutine_checker.is_async_callable(callback): - litellm._async_input_callback.append(callback) - removed_async_items.append(index) - - # Pop the async items from input_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.input_callback.pop(index) - if len(litellm.success_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.success_callback): - if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - removed_async_items.append(index) - elif callback == "dynamodb" or callback == "openmeter": - # dynamo is an async callback, it's used for the proxy and needs to be async - # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - removed_async_items.append(index) - elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): - add_custom_logger_callback_to_specific_event(callback, "success") - - # Pop the async items from success_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.success_callback.pop(index) - - if len(litellm.failure_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.failure_callback): - if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - removed_async_items.append(index) - elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): - add_custom_logger_callback_to_specific_event(callback, "failure") - - # Pop the async items from failure_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.failure_callback.pop(index) + get_set_callbacks()(callback_list=callbacks_to_set_up, function_id=function_id) ### DYNAMIC CALLBACKS ### - dynamic_success_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_async_success_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_async_failure_callbacks: Final[list[str | Callable | CustomLogger] | None] = None + dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = None if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list): removed_async_items = [] for index, callback in enumerate(kwargs["success_callback"]): @@ -1617,7 +1641,7 @@ def _deployment_provider(kwargs: Mapping[str, object]) -> str: return "" -async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str) -> dict[str, Any]: +async def async_pre_call_deployment_hook(kwargs: Mapping[str, object], call_type: str) -> dict[str, Any]: """ Allow modifying the request just before it's sent to the deployment. @@ -2597,19 +2621,14 @@ def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse: try: result: Final = _return_huggingface_tokenizer(model) - if result is not None: + if result is not None and not isinstance(result, _HuggingFaceTokenizerLoadFailure): return result except Exception as e: from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking if isinstance(e, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): raise - verbose_logger.warning( - "Falling back to tiktoken for %s; token counts may be approximate. " - "For Python Hugging Face tokenization, install tokenizers and huggingface-hub. Error: %s", - json.dumps(redact_secrets(model)), - json.dumps(redact_secrets(str(e))), - ) + _warn_huggingface_tokenizer_fallback(model, str(e)) # default - tiktoken return _return_openai_tokenizer(model) @@ -2649,15 +2668,41 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None: return None -def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None: +@dataclass(frozen=True, slots=True) +class _HuggingFaceTokenizerLoadFailure: + reason: str + + +def _warn_huggingface_tokenizer_fallback(model: str, reason: str) -> None: + verbose_logger.warning( + "Falling back to tiktoken for %s; token counts may be approximate. " + "For Python Hugging Face tokenization, install tokenizers and huggingface-hub. Error: %s", + json.dumps(redact_secrets(model)), + json.dumps(redact_secrets(reason)), + ) + + +def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _HuggingFaceTokenizerLoadFailure | None: kind: Final = huggingface_tokenizer_kind(model) - if kind is None: + if kind is None or (kind not in _BUNDLED_TOKENIZER_LOAD_LOCKS and is_offline_mode()): return None + backend: Final = _huggingface_tokenizer_backend() + key: Final = _huggingface_tokenizer_failure_key(kind, backend) with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()): - tokenizer: Final = _load_huggingface_tokenizer(kind, _huggingface_tokenizer_backend()) + cached_failure: Final = _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.get_cache(key) + if isinstance(cached_failure, _HuggingFaceTokenizerLoadFailure): + return cached_failure + tokenizer: Final = _load_huggingface_tokenizer_recording_failure(kind, backend, key) + if isinstance(tokenizer, _HuggingFaceTokenizerLoadFailure): + _warn_huggingface_tokenizer_fallback(model, tokenizer.reason) + return tokenizer return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} +def _huggingface_tokenizer_failure_key(kind: HuggingFaceTokenizerKind, backend: Decision) -> str: + return f"{kind}:{backend.name}" + + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decision) -> HuggingFace: """One tokenizer per kind and backend; `backend` is the cache key, the dispatch re-derives it.""" @@ -2680,6 +2725,21 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio assert_never(kind) +def _load_huggingface_tokenizer_recording_failure( + kind: HuggingFaceTokenizerKind, backend: Decision, key: str +) -> HuggingFace | _HuggingFaceTokenizerLoadFailure: + try: + return _load_huggingface_tokenizer(kind, backend) + except Exception as error: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking + + if isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): + raise + failure: Final = _HuggingFaceTokenizerLoadFailure(reason=str(error)) + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.set_cache(key, failure, ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS) + return failure + + def encode(model="", text="", custom_tokenizer: dict | None = None): """ Encodes the given text using the specified model. @@ -2904,12 +2964,14 @@ def supports_native_streaming(model: str, custom_llm_provider: str | None) -> bo Raises: Exception: If the given model is not found in model_prices_and_context_window.json. """ + from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata + try: - model, custom_llm_provider, _, _ = litellm.get_llm_provider( + resolved_model, resolved_provider = resolve_model_and_provider_for_metadata( model=model, custom_llm_provider=custom_llm_provider ) - model_info: Final = get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + model_info: Final = get_model_info_helper(model=resolved_model, custom_llm_provider=resolved_provider) supports_native_streaming = model_info.get("supports_native_streaming", True) if supports_native_streaming is None: supports_native_streaming = True @@ -4227,7 +4289,7 @@ def get_optional_params_embeddings( # retrieve all parameters passed to the function passed_params: Final = locals() - custom_llm_provider = passed_params.pop("custom_llm_provider", None) + passed_params.pop("custom_llm_provider", None) special_params: Final = passed_params.pop("kwargs") passed_params.pop("drop_params", None) @@ -4265,7 +4327,7 @@ def get_optional_params_embeddings( provider_config: BaseEmbeddingConfig | None = None optional_params = {} - if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values(): + if custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_embedding_config( model=model, provider=LlmProviders(custom_llm_provider), @@ -6713,6 +6775,13 @@ def get_model_info_helper( ), cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None), cache_read_input_audio_token_cost=_model_info.get("cache_read_input_audio_token_cost", None), + cache_read_input_audio_token_cost_flex=_model_info.get("cache_read_input_audio_token_cost_flex", None), + cache_read_input_audio_token_cost_priority=_model_info.get( + "cache_read_input_audio_token_cost_priority", None + ), + cache_read_input_audio_token_cost_batches=_model_info.get( + "cache_read_input_audio_token_cost_batches", None + ), explicit_cache_read_input_token_cost=_model_info.get("explicit_cache_read_input_token_cost", None), prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( @@ -6757,6 +6826,36 @@ def get_model_info_helper( cache_creation_input_token_cost_above_1hr=_model_info.get( "cache_creation_input_token_cost_above_1hr", None ), + cache_creation_input_token_cost_above_12hr=_model_info.get( + "cache_creation_input_token_cost_above_12hr", None + ), + cache_creation_input_token_cost_above_12hr_flex=_model_info.get( + "cache_creation_input_token_cost_above_12hr_flex", None + ), + cache_creation_input_token_cost_above_12hr_priority=_model_info.get( + "cache_creation_input_token_cost_above_12hr_priority", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority", None + ), + cache_creation_input_token_cost_above_12hr_batches=_model_info.get( + "cache_creation_input_token_cost_above_12hr_batches", None + ), + cache_creation_input_token_cost_above_12hr_ultrafast=_model_info.get( + "cache_creation_input_token_cost_above_12hr_ultrafast", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast", None + ), off_peak_pricing=_model_info.get("off_peak_pricing", None), input_cost_per_character=_model_info.get("input_cost_per_character", None), input_cost_per_token_above_128k_tokens=_model_info.get("input_cost_per_token_above_128k_tokens", None), @@ -6780,6 +6879,8 @@ def get_model_info_helper( input_cost_per_second=_model_info.get("input_cost_per_second", None), input_cost_per_text_message=_model_info.get("input_cost_per_text_message", None), input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None), + input_cost_per_audio_token_flex=_model_info.get("input_cost_per_audio_token_flex", None), + input_cost_per_audio_token_priority=_model_info.get("input_cost_per_audio_token_priority", None), input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None), input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None), input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None), @@ -7896,84 +7997,6 @@ def register_prompt_template( return litellm.custom_prompt_dict -class TextCompletionStreamWrapper: - def __init__( - self, - completion_stream, - model, - stream_options: dict | None = None, - custom_llm_provider: str | None = None, - ): - self.completion_stream = completion_stream - self.model = model - self.stream_options = stream_options - self.custom_llm_provider = custom_llm_provider - - def __iter__(self): - return self - - def __aiter__(self): - return self - - def convert_to_text_completion_object(self, chunk: ModelResponse): - try: - response: Final = TextCompletionResponse() - response["id"] = chunk.get("id", None) - response["object"] = "text_completion" - response["created"] = chunk.get("created", None) - response["model"] = chunk.get("model", None) - text_choices: Final = TextChoices() - if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices - raise Exception - delta: Final = chunk["choices"][0]["delta"] - text_choices["text"] = delta["content"] - text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] - response["choices"] = [text_choices] - - # only pass usage when stream_options["include_usage"] is True - if self.stream_options and self.stream_options.get("include_usage", False) is True: - response["usage"] = chunk.get("usage", None) - - return response - except Exception as e: - raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {e}") - - def __next__(self): - # model_response = ModelResponse(stream=True, model=self.model) - TextCompletionResponse() - try: - for chunk in self.completion_stream: - if chunk == "None" or chunk is None: - raise Exception - processed_chunk = self.convert_to_text_completion_object(chunk=chunk) - return processed_chunk - raise StopIteration - except StopIteration: - raise StopIteration - except Exception as e: - exception_type: Final = litellm_utils.exception_type - raise exception_type( - model=self.model, - custom_llm_provider=self.custom_llm_provider or "", - original_exception=e, - completion_kwargs={}, - extra_kwargs={}, - ) - - async def __anext__(self): - try: - async for chunk in self.completion_stream: - if chunk == "None" or chunk is None: - raise Exception - processed_chunk = self.convert_to_text_completion_object(chunk=chunk) - return processed_chunk - raise StopIteration - except StopIteration: - raise StopAsyncIteration - - def mock_stream_usage_chunk(model_response: ModelResponseStream, model: str, prompt_tokens: int) -> ModelResponseStream: return ModelResponseStream( id=model_response.id, @@ -8295,6 +8318,26 @@ def trim_messages( from litellm.caching.in_memory_cache import InMemoryCache +class _HuggingFaceTokenizerLoadFailureCache(Protocol): + def get_cache(self, key: str) -> object | None: ... + + def set_cache(self, key: str, value: object, ttl: float | None = None) -> None: ... + + def flush_cache(self) -> None: ... + + +_HUGGINGFACE_TOKENIZER_LOAD_FAILURES: Final[_HuggingFaceTokenizerLoadFailureCache] = ( + cast( # cast-ok: cache methods are untyped + _HuggingFaceTokenizerLoadFailureCache, + InMemoryCache( + max_size_in_memory=16, + default_ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, + clock=lambda: time.monotonic(), + ), + ) +) + + class AvailableModelsCache(InMemoryCache): def __init__(self, ttl_seconds: int = 300, max_size: int = 1000): super().__init__(ttl_seconds, max_size) @@ -9031,8 +9074,9 @@ _get_bundled_model_cost_map = get_bundled_model_cost_map def _get_model_cost_entry_for_provider_config( model: str, provider: LlmProviders, -) -> dict[str, object]: +) -> Mapping[str, object]: candidate_keys: Final = (model, f"{provider.value}/{model}") + model_info: Mapping[str, object] | None for model_key in candidate_keys: model_info = litellm.model_cost.get(model_key) if model_info is not None: @@ -9096,6 +9140,7 @@ class ProviderConfigManager: LlmProviders.DATABRICKS: (lambda: litellm.DatabricksConfig(), False), LlmProviders.XAI: (lambda: litellm.XAIChatConfig(), False), LlmProviders.ZAI: (lambda: litellm.ZAIChatConfig(), False), + LlmProviders.OPENCODE_GO: (ProviderConfigManager.get_opencode_go_config, True), LlmProviders.LAMBDA_AI: (lambda: litellm.LambdaAIChatConfig(), False), LlmProviders.INCEPTION: (lambda: litellm.InceptionChatConfig(), False), LlmProviders.LLAMA: (lambda: litellm.LlamaAPIConfig(), False), @@ -9287,6 +9332,12 @@ class ProviderConfigManager: return get_bedrock_chat_config(model=model) + @staticmethod + def get_opencode_go_config(model: str) -> BaseConfig: + from litellm.llms.opencode.chat.transformation import go_chat_config + + return go_chat_config(model) + @staticmethod def get_bedrock_mantle_config(model: str) -> BaseConfig: from litellm.llms.bedrock_mantle.chat.claude_transformation import bedrock_mantle_chat_config @@ -9522,7 +9573,7 @@ class ProviderConfigManager: return get_nvidia_nim_rerank_config(model) elif litellm.LlmProviders.VERTEX_AI == provider: - return litellm.VertexAIRerankConfig() + return ProviderConfigManager._get_vertex_ai_rerank_config_cached() elif litellm.LlmProviders.FIREWORKS_AI == provider: return litellm.FireworksAIRerankConfig() elif litellm.LlmProviders.VOYAGE == provider: @@ -9541,6 +9592,11 @@ class ProviderConfigManager: return get_dashscope_family_rerank_config(provider.value) return litellm.CohereRerankConfig() + @staticmethod + @lru_cache(maxsize=1) + def _get_vertex_ai_rerank_config_cached() -> BaseRerankConfig: + return litellm.VertexAIRerankConfig() + @staticmethod def get_provider_decisions_config(model: str, provider: LlmProviders) -> BaseDecisionsConfig | None: if provider == LlmProviders.PERPLEXITY: @@ -9590,6 +9646,11 @@ class ProviderConfigManager: model: str, provider: LlmProviders, ) -> BaseAnthropicMessagesConfig | None: + if provider == LlmProviders.OPENCODE_GO: + from litellm.llms.opencode.common_utils import go_endpoint + from litellm.llms.opencode.messages.transformation import OpenCodeGoMessagesConfig + + return OpenCodeGoMessagesConfig() if go_endpoint(model) == "/v1/messages" else None model_lower: Final = model.lower() if litellm.LlmProviders.ANTHROPIC == provider: return litellm.AnthropicMessagesConfig() @@ -9816,6 +9877,11 @@ class ProviderConfigManager: if provider is None: return None + if provider == LlmProviders.OPENCODE_GO: + from litellm.llms.opencode.common_utils import go_endpoint + from litellm.llms.opencode.responses.transformation import OpenCodeGoResponsesConfig + + return OpenCodeGoResponsesConfig() if go_endpoint(model or "") == "/v1/responses" else None if litellm.LlmProviders.OPENAI == provider: return litellm.OpenAIResponsesAPIConfig() elif litellm.LlmProviders.AZURE == provider: @@ -9955,6 +10021,10 @@ class ProviderConfigManager: model: str | None, provider: LlmProviders, ) -> BaseLLMModelInfo | None: + if provider == LlmProviders.OPENCODE_GO: + from litellm.llms.opencode.chat.transformation import OpenCodeGoChatConfig + + return OpenCodeGoChatConfig() if LlmProviders.FIREWORKS_AI == provider: return litellm.FireworksAIConfig() elif LlmProviders.BASETEN == provider: @@ -10857,18 +10927,12 @@ def extract_duration_from_srt_or_vtt(srt_or_vtt_content: str) -> float | None: # Regular expression to match timestamps in the format "hh:mm:ss,ms" or "hh:mm:ss.ms" timestamp_pattern: Final = r"(\d{2}):(\d{2}):(\d{2})[.,](\d{3})" - timestamps: Final[Sequence[tuple[str, str, str, str]]] = re.findall(timestamp_pattern, srt_or_vtt_content) - - if not timestamps: - return None + timestamps: Final = re.finditer(timestamp_pattern, srt_or_vtt_content) # Convert timestamps to seconds and find the max (end time) - durations: Final = [] - match: tuple[str, str, str, str] - for match in timestamps: - hours, minutes, seconds, milliseconds = map(int, match) - total_seconds = hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0 - durations.append(total_seconds) + durations: Final = [ + int(match[1]) * 3600 + int(match[2]) * 60 + int(match[3]) + int(match[4]) / 1000.0 for match in timestamps + ] return max(durations) if durations else None diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 8b33b1ce12c..9f53accfa0b 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -5,7 +5,7 @@ LiteLLM SDK Functions for Creating and Searching Vector Stores import asyncio import builtins import contextvars -from collections.abc import Coroutine, Mapping +from collections.abc import Coroutine, Iterator, Mapping from functools import partial from typing import TYPE_CHECKING, Final @@ -53,6 +53,10 @@ def _direct_vector_store_embedding_executor( ) +def _items_without_key(values: Mapping[str, object], excluded_key: str) -> Iterator[tuple[str, object]]: + return ((key, value) for key, value in values.items() if key != excluded_key) + + def mock_vector_store_search_response( mock_results: builtins.list[VectorStoreSearchResult] | None = None, ): @@ -307,7 +311,7 @@ async def asearch( embedding_executor: Final = _direct_vector_store_embedding_executor( kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs ) - local_vars: Final = {key: value for key, value in locals().items() if key != "embedding_executor"} + local_vars: Final = dict(_items_without_key(locals(), "embedding_executor")) try: loop: Final = asyncio.get_event_loop() @@ -391,7 +395,7 @@ def search( embedding_executor: Final = _direct_vector_store_embedding_executor( kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs ) - local_vars: Final = {key: value for key, value in locals().items() if key != "embedding_executor"} + local_vars: Final = dict(_items_without_key(locals(), "embedding_executor")) try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) @@ -694,7 +698,7 @@ def list( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -): +) -> Mapping[str, object] | Coroutine[object, object, Mapping[str, object]]: """ List vector stores. @@ -992,7 +996,7 @@ def delete( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -): +) -> Mapping[str, object] | Coroutine[object, object, Mapping[str, object]]: """ Delete a vector store. diff --git a/migrations/Dockerfile b/migrations/Dockerfile index b39a9bfa9bf..2b1fd029143 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -61,7 +61,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/ # instead of nodeenv downloading one whose dynamic deps may not be in Wolfi # (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -118,7 +118,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python-3.13 nodejs libsndfile libatomic && break; \ + apk add --no-cache bash openssl libssl3 tzdata python-3.13 nodejs libsndfile libatomic && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -137,7 +137,9 @@ ENV HOME=/home/nonroot \ PRISMA_CLI_QUERY_ENGINE_TYPE=binary \ PRISMA_OFFLINE_MODE=true -COPY --from=builder --chown=nonroot:nonroot /app /app +COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv +COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm +COPY --from=builder --chown=nonroot:nonroot /app/run.py /app/run.py COPY --from=builder /opt/prisma /opt/prisma RUN chmod -R a+rX /opt/prisma && \ diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3d89af9d691..d9df9d97d7b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -644,9 +644,13 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 3.3e-07, "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2.75e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-nova-2-sonic.md", "supports_audio_input": true, "supports_audio_output": true }, @@ -4339,7 +4343,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -6202,6 +6206,8 @@ "text", "audio" ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6266,6 +6272,8 @@ "text", "audio" ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -9232,7 +9240,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -9679,7 +9687,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5.5e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -12827,6 +12835,7 @@ "supports_vision": true }, "azure_ai/Phi-4": { + "deprecation_date": "2027-01-17", "input_cost_per_token": 1.25e-07, "litellm_provider": "azure_ai", "max_input_tokens": 16384, @@ -12840,6 +12849,7 @@ "supports_vision": false }, "azure_ai/Phi-4-mini-instruct": { + "deprecation_date": "2027-01-17", "input_cost_per_token": 7.5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12851,6 +12861,7 @@ "supports_function_calling": true }, "azure_ai/Phi-4-multimodal-instruct": { + "deprecation_date": "2027-01-17", "input_cost_per_audio_token": 4e-06, "input_cost_per_token": 8e-08, "litellm_provider": "azure_ai", @@ -12865,6 +12876,7 @@ "supports_vision": true }, "azure_ai/Phi-4-mini-reasoning": { + "deprecation_date": "2027-01-17", "input_cost_per_token": 7.5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12876,6 +12888,7 @@ "supports_function_calling": true }, "azure_ai/Phi-4-reasoning": { + "deprecation_date": "2027-01-17", "input_cost_per_token": 1.25e-07, "litellm_provider": "azure_ai", "max_input_tokens": 32768, @@ -15654,6 +15667,28 @@ "/v1/images/generations" ] }, + "cerebras/llama-3.1-70b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "supports_function_calling": true, + "supports_tool_choice": true + }, + "cerebras/llama-3.1-8b": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-07, + "supports_function_calling": true, + "supports_tool_choice": true + }, "cerebras/llama-3.3-70b": { "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", @@ -24892,27 +24927,118 @@ ] }, "elevenlabs/eleven_v3": { - "input_cost_per_character": 0.00018, + "input_cost_per_character": 8e-05, "litellm_provider": "elevenlabs", "metadata": { - "calculation": "$0.18/1000 characters (Scale plan pricing, 1 credit per character)", + "calculation": "$0.08/1000 characters (API usage pricing, billed in USD per character)", "notes": "ElevenLabs Eleven v3 - most expressive TTS model with 70+ languages and audio tags support" }, "mode": "audio_speech", - "source": "https://elevenlabs.io/pricing", + "source": "https://elevenlabs.io/pricing/api", "supported_endpoints": [ "/v1/audio/speech" ] }, "elevenlabs/eleven_multilingual_v2": { - "input_cost_per_character": 0.00018, + "input_cost_per_character": 8e-05, "litellm_provider": "elevenlabs", "metadata": { - "calculation": "$0.18/1000 characters (Scale plan pricing, 1 credit per character)", + "calculation": "$0.08/1000 characters (API usage pricing, billed in USD per character)", "notes": "ElevenLabs Eleven Multilingual v2 - default TTS model with 29 languages support" }, "mode": "audio_speech", - "source": "https://elevenlabs.io/pricing", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_v4": { + "input_cost_per_character": 8e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.08/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven v4 - most emotive TTS model with audio tags and 90+ languages" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_v4_turbo": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven v4 Turbo - low-latency v4 tuned for real-time and agents" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_v3_conversational": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven v3 Conversational - low-latency v3 tuned for realtime conversation" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_flash_v2_5": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven Flash v2.5 - low-latency multilingual TTS model" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_flash_v2": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven Flash v2 - low-latency English-only TTS model" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_turbo_v2_5": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven Turbo v2.5 - multilingual TTS model" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "elevenlabs/eleven_turbo_v2": { + "input_cost_per_character": 4e-05, + "litellm_provider": "elevenlabs", + "metadata": { + "calculation": "$0.04/1000 characters (API usage pricing, billed in USD per character)", + "notes": "ElevenLabs Eleven Turbo v2 - English-only TTS model" + }, + "mode": "audio_speech", + "source": "https://elevenlabs.io/pricing/api", "supported_endpoints": [ "/v1/audio/speech" ] @@ -31041,11 +31167,15 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_batches": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31370,11 +31500,15 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_audio_token_cost": 3e-08, + "cache_read_input_audio_token_cost_batches": 3e-08, + "cache_read_input_audio_token_cost_flex": 3e-08, + "cache_read_input_audio_token_cost_priority": 5.4e-08, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_audio_token": 3e-07, + "input_cost_per_audio_token_flex": 1.5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31664,12 +31798,16 @@ }, "gemini/gemini-3.1-flash-lite": { "cache_read_input_audio_token_cost": 5e-08, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "cache_read_input_audio_token_cost_priority": 9e-08, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "deprecation_date": "2027-05-07", "input_cost_per_audio_token": 5e-07, + "input_cost_per_audio_token_flex": 2.5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -31790,10 +31928,14 @@ }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 5e-08, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -33100,6 +33242,7 @@ "github_copilot/claude-opus-4.8-fast": { "cache_creation_input_token_cost": 1.25e-05, "cache_read_input_token_cost": 1e-06, + "deprecation_date": "2026-11-16", "input_cost_per_token": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 200000, @@ -33559,7 +33702,7 @@ "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33583,7 +33726,7 @@ "input_cost_per_token": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33607,7 +33750,7 @@ "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33631,7 +33774,7 @@ "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -36699,6 +36842,16 @@ }, "gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_12hr": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, + "cache_creation_input_token_cost_above_12hr_priority": 5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -36784,6 +36937,14 @@ }, "gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, @@ -36861,17 +37022,25 @@ }, "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, - "cache_creation_input_token_cost_decisions": 1e-07, + "cache_creation_input_token_cost_above_12hr": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, + "cache_creation_input_token_cost_above_12hr_priority": 5e-07, + "cache_creation_input_token_cost_decisions": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, - "cache_creation_input_token_cost_above_272k_tokens_decisions": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_decisions": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-07, "cache_creation_input_token_cost_flex": 6.25e-08, "cache_creation_input_token_cost_priority": 2.5e-07, "cache_read_input_token_cost": 1e-08, - "cache_read_input_token_cost_decisions": 1e-07, + "cache_read_input_token_cost_decisions": 1e-08, "cache_read_input_token_cost_above_272k_tokens": 2e-08, - "cache_read_input_token_cost_above_272k_tokens_decisions": 2e-07, + "cache_read_input_token_cost_above_272k_tokens_decisions": 2e-08, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 4e-08, "cache_read_input_token_cost_batches": 5e-09, @@ -36947,6 +37116,12 @@ }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_12hr": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, @@ -37021,6 +37196,14 @@ }, "gpt-5.6-sol": { "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_12hr": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_batches": 5e-06, + "cache_creation_input_token_cost_above_12hr_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, @@ -37097,6 +37280,14 @@ }, "gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, @@ -37172,6 +37363,14 @@ }, "gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_12hr": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_priority": 1e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, @@ -37247,6 +37446,7 @@ }, "gpt-5.6-cyber": { "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_above_12hr": 3.125e-05, "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, @@ -42337,6 +42537,134 @@ ], "supports_audio_output": true }, + "mistral/voxtral-mini-transcribe-3": { + "input_cost_per_second": 5.83333e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-mini-transcribe-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-transcribe-3-0": { + "input_cost_per_second": 5.83333e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-mini-transcribe-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-small-transcribe-3": { + "input_cost_per_second": 6.66667e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-small-transcribe-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-small-transcribe-3-0": { + "input_cost_per_second": 6.66667e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-small-transcribe-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-transcribe-realtime-3": { + "input_cost_per_second": 0.000133333, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-mini-transcribe-realtime-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-transcribe-realtime-3-0": { + "input_cost_per_second": 0.000133333, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/voxtral-mini-transcribe-realtime-3-0", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-tts-3": { + "input_cost_per_character": 1.3e-05, + "litellm_provider": "mistral", + "mode": "audio_speech", + "source": "https://docs.mistral.ai/models/voxtral-mini-tts-3-0", + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_output": true + }, + "mistral/voxtral-mini-tts-3-0": { + "input_cost_per_character": 1.3e-05, + "litellm_provider": "mistral", + "mode": "audio_speech", + "source": "https://docs.mistral.ai/models/voxtral-mini-tts-3-0", + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_output": true + }, "mistral/voxtral-small-2507": { "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 1e-07, @@ -42416,6 +42744,29 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5e-07, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://docs.mistral.ai/models/zai-glm-5-3-flash", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/zai-glm-5": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -45623,6 +45974,7 @@ "output_cost_per_token": 0.0 }, "openai.gpt-oss-120b-1:0": { + "supports_vision": false, "supported_endpoints": [ "/v1/chat/completions" ], @@ -45648,6 +46000,7 @@ "supports_tool_choice": true }, "openai.gpt-oss-20b-1:0": { + "supports_vision": false, "supported_endpoints": [ "/v1/chat/completions" ], @@ -48499,6 +48852,17 @@ "supports_tool_choice": true, "supports_vision": true }, + "together_ai/MiniMaxAI/MiniMax-M3.1-Flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://www.together.ai/pricing", + "supports_prompt_caching": true + }, "together_ai/Prism-ML/Ternary-Bonsai-27B": { "input_cost_per_token": 0.0, "litellm_provider": "together_ai", @@ -49352,6 +49716,7 @@ "supports_vision": false }, "us-gov.openai.gpt-oss-20b-1:0": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, @@ -49371,6 +49736,7 @@ "supports_tool_choice": true }, "us-gov.openai.gpt-oss-120b-1:0": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, @@ -52906,12 +53272,14 @@ }, "vertex_ai/meta/llama-4-maverick-17b-16e-instruct-maas": { "input_cost_per_token": 3.5e-07, + "input_cost_per_token_batches": 1.75e-07, "litellm_provider": "vertex_ai-llama_models", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 1.15e-06, + "output_cost_per_token_batches": 5.75e-07, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", "supported_modalities": [ "text", @@ -52926,12 +53294,14 @@ }, "vertex_ai/meta/llama-4-scout-17b-128e-instruct-maas": { "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, "litellm_provider": "vertex_ai-llama_models", "max_input_tokens": 10000000, "max_output_tokens": 10000000, "max_tokens": 10000000, "mode": "chat", "output_cost_per_token": 7e-07, + "output_cost_per_token_batches": 3.5e-07, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", "supported_modalities": [ "text", @@ -59213,6 +59583,7 @@ "supports_audio_input": true, "supports_audio_output": true, "gemini_native_audio": true, + "input_cost_per_video_token": 3e-06, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -59408,6 +59779,7 @@ "tpm": 250000, "rpm": 10, "gemini_native_audio": true, + "input_cost_per_video_token": 3e-06, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -59917,6 +60289,7 @@ } }, "bedrock_mantle/openai.gpt-oss-120b": { + "supports_vision": false, "input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, @@ -59942,6 +60315,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-20b": { + "supports_vision": false, "input_cost_per_token": 7e-08, "input_cost_per_token_batches": 3.5e-08, "input_cost_per_token_flex": 3.5e-08, @@ -59967,6 +60341,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-120b": { + "supports_vision": false, "input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, @@ -59990,6 +60365,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { + "supports_vision": false, "input_cost_per_token": 7e-08, "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, @@ -64195,13 +64571,17 @@ }, "fireworks_ai/accounts/fireworks/models/muse-spark-open": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_priority": 2.25e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_priority": 1.875e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "mode": "chat", "output_cost_per_token": 4.25e-06, - "source": "https://fireworks.ai/models/fireworks/muse-spark-open", + "output_cost_per_token_priority": 6.375e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_reasoning": true, "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": { @@ -67837,6 +68217,7 @@ "supports_vision": false }, "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -67856,6 +68237,7 @@ "supports_tool_choice": true }, "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -68134,6 +68516,7 @@ "supports_vision": false }, "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -68153,6 +68536,7 @@ "supports_tool_choice": true }, "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -68541,6 +68925,7 @@ "supports_vision": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", @@ -68560,6 +68945,7 @@ "supports_tool_choice": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", @@ -68700,6 +69086,7 @@ "supports_vision": true }, "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": { + "supports_vision": false, "input_cost_per_token": 8.4e-08, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", @@ -68719,6 +69106,7 @@ "supports_tool_choice": true }, "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": { + "supports_vision": false, "input_cost_per_token": 1.8e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", @@ -73326,6 +73714,7 @@ ] }, "vertex_ai/veo-2.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -74673,6 +75062,16 @@ }, "gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, @@ -75038,10 +75437,13 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 3.3e-07, "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2.75e-06, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-nova-2-5-sonic.md", "supports_audio_input": true, "supports_audio_output": true }, @@ -76059,9 +76461,10 @@ "azure_ai/Microsoft-Decision-1": { "input_cost_per_token": 4.2e-08, "litellm_provider": "azure_ai", + "max_input_tokens": 32768, "mode": "evaluation", "output_cost_per_token": 0.0, - "source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'" + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" }, "gpt-rosalind-discovery": { "cache_read_input_token_cost": 5e-07, @@ -76709,14 +77112,14 @@ "supports_response_schema": true }, "baseten/reflection/Beam-1.0": { - "cache_read_input_token_cost": 1.3e-07, - "input_cost_per_token": 1.3e-06, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, "litellm_provider": "baseten", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4.3e-06, + "output_cost_per_token": 1.2e-06, "source": "https://inference.baseten.co/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -77018,5 +77421,1481 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false + }, + "azure_ai/MAI-Code-1.1-Flash": { + "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2026-12-09", + "input_cost_per_token": 2e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 256000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "opencode_go/grok-4.7": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 5e-07, + "max_tokens": 500000, + "max_output_tokens": 500000, + "max_input_tokens": 500000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 200000 + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 5e-07 + }, + { + "range": [ + 200000, + 500000 + ], + "input_cost_per_token": 4e-06, + "output_cost_per_token": 1.2e-05, + "cache_read_input_token_cost": 1e-06 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/gpt-6-luna": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_creation_input_token_cost": 1.25e-07, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 922000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 272000 + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_creation_input_token_cost": 1.25e-07 + }, + { + "range": [ + 272000, + 1050000 + ], + "input_cost_per_token": 2.0000000000000002e-07, + "output_cost_per_token": 7.5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_creation_input_token_cost": 2.5e-07 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/gpt-5.6-luna": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 2.0000000000000002e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 2e-08, + "cache_creation_input_token_cost": 2.5e-07, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 922000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 272000 + ], + "input_cost_per_token": 2.0000000000000002e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 2e-08, + "cache_creation_input_token_cost": 2.5e-07 + }, + { + "range": [ + 272000, + 1050000 + ], + "input_cost_per_token": 4.0000000000000003e-07, + "output_cost_per_token": 1.8000000000000001e-06, + "cache_read_input_token_cost": 4e-08, + "cache_creation_input_token_cost": 5e-07 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/claude-haiku-5-5": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_creation_input_token_cost": 1.25e-07, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 100000 + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "cache_creation_input_token_cost": 1.25e-07 + }, + { + "range": [ + 100000, + 1000000 + ], + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 5.0000000000000004e-08, + "cache_creation_input_token_cost": 6.25e-07 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/gemini-4-argon": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1beta/models" + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1e-07, + "max_tokens": 65536, + "max_output_tokens": 65536, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/glm-5.3-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 3e-08, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/glm-5.3": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/glm-5.2": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/kimi-k3.1-pro": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2.0000000000000002e-07, + "cache_creation_input_token_cost": 2e-06, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/kimi-k3": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3e-06, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/kimi-k2.7-code": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 9.499999999999999e-07, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 1.9e-07, + "max_tokens": 262144, + "max_output_tokens": 262144, + "max_input_tokens": 262144, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/longcat-2.0": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/deepseek-v4.1-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-09, + "max_tokens": 384000, + "max_output_tokens": 384000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/deepseek-v4-pro": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.32e-06, + "output_cost_per_token": 3.96e-06, + "cache_read_input_token_cost": 4.4e-08, + "max_tokens": 384000, + "max_output_tokens": 384000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/mimo-v2.6-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 2.8e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/mimo-v2.6-pro": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 4.35e-07, + "output_cost_per_token": 8.7e-07, + "cache_read_input_token_cost": 3.625e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/mimo-v2.5": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 2.8e-09, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/mimo-v2.5-pro": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 4.35e-07, + "output_cost_per_token": 8.7e-07, + "cache_read_input_token_cost": 3.625e-09, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/minimax-m3.1-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 1.5e-08, + "max_tokens": 512000, + "max_output_tokens": 512000, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/minimax-m3": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-08, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/muse-spark-1.3-contributor": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 2.0000000000000002e-07, + "cache_read_input_token_cost": 2e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/muse-spark-1.2-contributor": { + "litellm_provider": "opencode_go", + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "input_cost_per_token": 1.0000000000000001e-07, + "output_cost_per_token": 2.0000000000000002e-07, + "cache_read_input_token_cost": 2e-09, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/qwen3.8-max": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost": 2.5e-06, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/qwen3.8-flash": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 4.6999999999999995e-07, + "cache_read_input_token_cost": 1.6e-08, + "cache_creation_input_token_cost": 2.0000000000000002e-07, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/qwen3.7-plus": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/messages" + ], + "input_cost_per_token": 4.0000000000000003e-07, + "output_cost_per_token": 1.6000000000000001e-06, + "cache_read_input_token_cost": 4e-08, + "cache_creation_input_token_cost": 5e-07, + "max_tokens": 65536, + "max_output_tokens": 65536, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "tiered_pricing": [ + { + "range": [ + 0, + 256000 + ], + "input_cost_per_token": 4.0000000000000003e-07, + "output_cost_per_token": 1.6000000000000001e-06, + "cache_read_input_token_cost": 4e-08, + "cache_creation_input_token_cost": 5e-07 + }, + { + "range": [ + 256000, + 1000000 + ], + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 4.8e-06, + "cache_read_input_token_cost": 1.2e-07, + "cache_creation_input_token_cost": 1.5e-06 + } + ], + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/hy4-preview": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 8.34e-07, + "output_cost_per_token": 2.501e-06, + "cache_read_input_token_cost": 4.2000000000000006e-08, + "max_tokens": 64000, + "max_output_tokens": 64000, + "max_input_tokens": 1024000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/hy3": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 5.8e-07, + "cache_read_input_token_cost": 3.5e-08, + "max_tokens": 128000, + "max_output_tokens": 128000, + "max_input_tokens": 192000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/longcat-2.5-preview-free": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "cache_read_input_token_cost": 0, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/step-5-preview": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 1.0000000000000001e-07, + "max_tokens": 65536, + "max_output_tokens": 65536, + "max_input_tokens": 1000000, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "opencode_go/beam-free": { + "litellm_provider": "opencode_go", + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "cache_read_input_token_cost": 0, + "max_tokens": 131072, + "max_output_tokens": 131072, + "max_input_tokens": 1048576, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false, + "source": "https://opencode.ai/v2/docs/console/go/" + }, + "bedrock/claude_platform/claude-fable-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-fable-5-1": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_computer_use": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-5-5": { + "input_cost_per_token_above_100k_tokens_batches": 2.5e-07, + "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, + "cache_read_input_token_cost_above_100k_tokens_batches": 2.5e-08, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_1hr": 2e-07, + "cache_creation_input_token_cost_batches": 6.25e-08, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_batches": 5e-09, + "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_batches": 2.5e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": true, + "thinking_always_on": false, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true, + "input_cost_per_token_above_100k_tokens": 5e-07, + "output_cost_per_token_above_100k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, + "cache_read_input_token_cost_above_100k_tokens": 5e-08 + }, + "bedrock/claude_platform/claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-7": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 2048, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-8": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_creation_input_token_cost_batches": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost_batches": 1.5e-07, + "deprecation_date": "2026-11-30", + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "provider_specific_entry": { + "us": 1.1 + }, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5-5": { + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": false, + "thinking_always_on": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true + }, + "gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_12hr": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-07, + "cache_creation_input_token_cost_priority": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "cache_read_input_token_cost_above_272k_tokens_batches": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-07, + "cache_read_input_token_cost_batches": 2.5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, + "cache_read_input_token_cost_priority": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "input_cost_per_token_above_272k_tokens_batches": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_priority": 4e-06, + "input_cost_per_token_batches": 5e-07, + "input_cost_per_token_flex": 5e-07, + "input_cost_per_token_priority": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "output_cost_per_token_above_272k_tokens_batches": 3.75e-06, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.5e-05, + "output_cost_per_token_batches": 2.5e-06, + "output_cost_per_token_flex": 2.5e-06, + "output_cost_per_token_priority": 1e-05, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6.1-terra": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "global.openai.gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_bedrock_converse_tool_result_images": false, + "supports_bedrock_runtime_chat_completions_response_format": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "us.openai.gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, + "cache_read_input_token_cost": 5.5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "input_cost_per_token_above_272k_tokens": 2.2e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "output_cost_per_token_above_272k_tokens": 8.25e-06, + "source": "https://developers.openai.com/api/docs/models/gpt-6.1-terra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_bedrock_converse_tool_result_images": false, + "supports_bedrock_runtime_chat_completions_response_format": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "baseten/decision-models/baseten-elephant-9b-latest": { + "input_cost_per_token": 0.0, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true + }, + "baseten/decision-models/baseten-mammoth-27b-latest": { + "input_cost_per_token": 0.0, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true + }, + "github_copilot/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json deleted file mode 100644 index d4dfeb25429..00000000000 --- a/model_prices_and_context_window.schema.json +++ /dev/null @@ -1,1312 +0,0 @@ -{ - "$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": { - "type": "string" - }, - "pattern": { - "type": "string" - }, - "description": { - "type": "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": { - "annotation_cost_per_page": { - "type": "number", - "minimum": 0 - }, - "annotation_cost_per_page_batches": { - "type": "number", - "minimum": 0 - }, - "audio_transcription_config": { - "type": "string" - }, - "bedrock_converse_supports_strict_tools": { - "type": "boolean" - }, - "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" - ] - }, - "cache_creation_input_audio_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost": { - "type": "number", - "minimum": 0, - "description": "USD per token written to the provider's prompt cache." - }, - "cache_creation_input_token_cost_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_100k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_1hr": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_1hr_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_1hr_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_272k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_272k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_272k_tokens_decisions": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_272k_tokens_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "cache_creation_input_token_cost_above_272k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_creation_input_token_cost_above_272k_tokens_ultrafast": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_batches": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost_decisions": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "cache_creation_input_token_cost_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_creation_input_token_cost_ultrafast": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_audio_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_image_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost": { - "type": "number", - "minimum": 0, - "description": "USD per prompt token served from the provider's prompt cache." - }, - "cache_read_input_token_cost_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_100k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_200k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_200k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_above_272k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_272k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_272k_tokens_decisions": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_272k_tokens_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_above_272k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_above_272k_tokens_ultrafast": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_512k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_balanced": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost_batches": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost_decisions": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_ultrafast": { - "type": "number", - "minimum": 0 - }, - "citation_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "code_interpreter_cost_per_session": { - "type": "number", - "minimum": 0 - }, - "comment": { - "type": "string" - }, - "cost_per_second": { - "type": "number", - "minimum": 0 - }, - "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" - ] - }, - "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])$" - }, - "explicit_cache_read_input_token_cost": { - "type": "number", - "minimum": 0 - }, - "gemini_audio_only_live": { - "type": "boolean" - }, - "gemini_native_audio": { - "type": "boolean" - }, - "google_maps_grounding_cost_per_query": { - "type": "number", - "minimum": 0, - "description": "USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit." - }, - "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": { - "type": "number", - "minimum": 0 - } - }, - "input_cost_per_audio_per_second": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_audio_per_second_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_audio_token": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_audio_token_batches": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_audio_token_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "input_cost_per_character": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_character_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_image": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_image_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_image_token": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_image_token_batches": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_pixel": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_query": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_request": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_second": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_text_message": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_token": { - "type": "number", - "minimum": 0, - "description": "USD per prompt token." - }, - "input_cost_per_token_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_100k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_200k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_200k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "input_cost_per_token_above_272k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_272k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_272k_tokens_decisions": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_272k_tokens_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "input_cost_per_token_above_272k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "input_cost_per_token_above_272k_tokens_ultrafast": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_512k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_balanced": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_token_batches": { - "type": "number", - "minimum": 0, - "description": "USD per prompt token via the provider's batch API." - }, - "input_cost_per_token_cache_hit": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_token_decisions": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_token_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "input_cost_per_token_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "input_cost_per_token_ultrafast": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_video_per_second": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_video_per_second_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_video_per_second_above_15s_interval": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_video_per_second_above_8s_interval": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_video_token": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_video_token_batches": { - "type": "number", - "minimum": 0 - }, - "input_dbu_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "litellm_provider": { - "type": "string", - "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers." - }, - "max_input_tokens": { - "type": "integer", - "minimum": 0, - "description": "Maximum prompt/context tokens the model accepts." - }, - "max_output_tokens": { - "type": "integer", - "minimum": 0, - "description": "Maximum tokens the model can generate in one response." - }, - "max_tokens": { - "type": "integer", - "minimum": 0, - "description": "Legacy field: max output tokens if the provider specifies it, else max input tokens." - }, - "metadata": { - "type": "object", - "description": "Free-form notes about the entry (e.g. pricing derivation)." - }, - "mode": { - "type": "string", - "description": "Primary API surface / task type of the model.", - "enum": [ - "audio_speech", - "audio_transcription", - "chat", - "completion", - "embedding", - "evaluation", - "extraction", - "guardrail", - "image_edit", - "image_generation", - "moderation", - "ocr", - "realtime", - "rerank", - "responses", - "search", - "vector_store", - "video_generation" - ] - }, - "ocr_cost_per_credit": { - "type": "number", - "minimum": 0 - }, - "ocr_cost_per_page": { - "type": "number", - "minimum": 0 - }, - "ocr_cost_per_page_batches": { - "type": "number", - "minimum": 0 - }, - "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": { - "description": "UTC \"HH:MM-HH:MM\" window, or a list of them; a window may wrap past midnight.", - "oneOf": [ - { - "type": "string", - "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" - }, - { - "type": "array", - "items": { - "type": "string", - "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" - }, - "minItems": 1 - } - ] - }, - "windows": { - "type": "array", - "items": { - "type": "object", - "properties": { - "hours_utc": { - "description": "UTC \"HH:MM-HH:MM\" window, or a list of them; a window may wrap past midnight.", - "oneOf": [ - { - "type": "string", - "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" - }, - { - "type": "array", - "items": { - "type": "string", - "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" - }, - "minItems": 1 - } - ] - }, - "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": "(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" - } - ] - }, - "minItems": 1 - } - }, - "required": [ - "hours_utc" - ], - "additionalProperties": false - }, - "minItems": 1 - }, - "weekday_timezone": { - "type": "string", - "description": "IANA zone the weekdays of each window are read on; defaults to UTC." - }, - "input_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_reasoning_token": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost": { - "type": "number", - "minimum": 0 - } - }, - "anyOf": [ - { - "required": [ - "hours_utc" - ] - }, - { - "required": [ - "windows" - ] - } - ], - "additionalProperties": false - }, - "output_cost_per_audio_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_character": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_character_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_image": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_0.5K": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_1024": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_1536": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_1K": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_2K": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_4K": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_512": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_token_batches": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_pixel": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_reasoning_token": { - "type": "number", - "minimum": 0, - "description": "USD per reasoning/thinking token, when billed separately." - }, - "output_cost_per_second": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_1080p": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_2k": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_480p": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_4k": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_720p": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_768p": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token": { - "type": "number", - "minimum": 0, - "description": "USD per generated token." - }, - "output_cost_per_token_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_100k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_200k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_200k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "output_cost_per_token_above_272k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_272k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_272k_tokens_decisions": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_272k_tokens_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "output_cost_per_token_above_272k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "output_cost_per_token_above_272k_tokens_ultrafast": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_512k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_balanced": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token_batches": { - "type": "number", - "minimum": 0, - "description": "USD per generated token via the provider's batch API." - }, - "output_cost_per_token_decisions": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "output_cost_per_token_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "output_cost_per_token_ultrafast": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_video_per_second": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_video_token": { - "type": "number", - "minimum": 0 - }, - "output_dbu_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_vector_size": { - "type": "integer", - "minimum": 0, - "description": "Embedding dimension for embedding models." - }, - "prompt_cache_min_tokens": { - "type": "integer", - "minimum": 0, - "description": "Smallest prefix the provider will actually cache; absent means the provider default applies." - }, - "provider_specific_entry": { - "type": "object", - "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema)." - }, - "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" - ] - } - }, - "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%)." - }, - "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%)." - }, - "rpm": { - "type": "integer", - "minimum": 0, - "description": "Provider default requests-per-minute limit." - }, - "search_context_cost_per_query": { - "type": "object", - "description": "USD cost per web search query, keyed by search context size.", - "properties": { - "search_context_size_low": { - "type": "number", - "minimum": 0 - }, - "search_context_size_medium": { - "type": "number", - "minimum": 0 - }, - "search_context_size_high": { - "type": "number", - "minimum": 0 - } - }, - "additionalProperties": false - }, - "search_free_requests": { - "type": "object", - "description": "Free Google Search grounding requests per period, shared across models in the same shared_group.", - "required": [ - "count", - "period", - "shared_group" - ], - "additionalProperties": false, - "properties": { - "count": { - "type": "integer", - "minimum": 0 - }, - "period": { - "type": "string", - "enum": [ - "day", - "month" - ] - }, - "shared_group": { - "type": "string", - "pattern": "\\S" - } - } - }, - "source": { - "type": "string", - "description": "URL of the provider pricing/model page this entry was taken from." - }, - "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": { - "type": "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" - ] - } - }, - "supported_regions": { - "type": "array", - "description": "Cloud regions the model is available in ('global' or region ids).", - "items": { - "type": "string" - } - }, - "supports_adaptive_thinking": { - "type": "boolean" - }, - "supports_anthropic_compaction": { - "type": "boolean" - }, - "supports_anthropic_thinking_payload": { - "type": "boolean" - }, - "supports_assistant_prefill": { - "type": "boolean" - }, - "supports_audio_input": { - "type": "boolean" - }, - "supports_audio_output": { - "type": "boolean" - }, - "supports_bedrock_converse_tool_result_images": { - "type": "boolean" - }, - "supports_bedrock_native_structured_output": { - "type": "boolean" - }, - "supports_bedrock_runtime_chat_completions_inline_reasoning": { - "type": "boolean" - }, - "supports_bedrock_runtime_chat_completions_response_format": { - "type": "boolean" - }, - "supports_bedrock_runtime_chat_completions_tools_with_reasoning": { - "type": "boolean" - }, - "supports_box_questions": { - "type": "boolean" - }, - "supports_computer_use": { - "type": "boolean" - }, - "supports_embedding_image_input": { - "type": "boolean" - }, - "supports_fast_mode": { - "type": "boolean" - }, - "supports_forced_tool_use": { - "type": "boolean" - }, - "supports_function_calling": { - "type": "boolean" - }, - "supports_image_input": { - "type": "boolean" - }, - "supports_image_size": { - "type": "boolean" - }, - "supports_legacy_thinking": { - "type": "boolean" - }, - "supports_low_reasoning_effort": { - "type": "boolean" - }, - "supports_max_reasoning_effort": { - "type": "boolean" - }, - "supports_mid_conversation_system": { - "type": "boolean" - }, - "supports_minimal_reasoning_effort": { - "type": "boolean" - }, - "supports_multimodal": { - "type": "boolean" - }, - "supports_native_streaming": { - "type": "boolean" - }, - "supports_native_structured_output": { - "type": "boolean" - }, - "supports_none_reasoning_effort": { - "type": "boolean" - }, - "supports_nova_canvas_image_edit": { - "type": "boolean" - }, - "supports_output_config": { - "type": "boolean" - }, - "supports_parallel_function_calling": { - "type": "boolean" - }, - "supports_parallel_tool_use_config": { - "type": "boolean" - }, - "supports_pdf_input": { - "type": "boolean" - }, - "supports_prompt_cache_breakpoint": { - "type": "boolean" - }, - "supports_prompt_caching": { - "type": "boolean" - }, - "supports_reasoning": { - "type": "boolean" - }, - "supports_regex_lookaround": { - "type": "boolean" - }, - "supports_response_schema": { - "type": "boolean" - }, - "supports_sampling_params": { - "type": "boolean" - }, - "supports_speed": { - "type": "boolean" - }, - "supports_system_messages": { - "type": "boolean" - }, - "supports_thinking_cache_preservation": { - "type": "boolean" - }, - "supports_tool_choice": { - "type": "boolean" - }, - "supports_tool_search": { - "type": "boolean" - }, - "supports_url_context": { - "type": "boolean" - }, - "supports_video_input": { - "type": "boolean" - }, - "supports_vision": { - "type": "boolean" - }, - "supports_web_search": { - "type": "boolean" - }, - "supports_xhigh_reasoning_effort": { - "type": "boolean" - }, - "thinking_always_on": { - "type": "boolean" - }, - "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": { - "type": "number", - "minimum": 0 - }, - "minItems": 2, - "maxItems": 2 - }, - "max_results_range": { - "type": "array", - "description": "[min, max] result-count span this tier applies to (search models).", - "items": { - "type": "number", - "minimum": 0 - }, - "minItems": 2, - "maxItems": 2 - }, - "input_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_reasoning_token": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_query": { - "type": "number", - "minimum": 0 - } - }, - "additionalProperties": false - } - }, - "token_count_framing": { - "type": "integer", - "minimum": 0, - "description": "Tokens added once per request to a local count estimate, for the provider's chat framing." - }, - "token_count_model": { - "type": "object", - "required": [ - "model" - ], - "additionalProperties": false, - "properties": { - "model": { - "type": "string", - "pattern": "\\S" - }, - "region": { - "type": "string", - "pattern": "\\S" - }, - "litellm_provider": { - "type": "string", - "enum": [ - "bedrock", - "bedrock_mantle" - ] - } - }, - "description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in." - }, - "token_count_tokenizer": { - "type": "string", - "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", - "enum": [ - "kimi_k3" - ] - }, - "tpm": { - "type": "integer", - "minimum": 0, - "description": "Provider default tokens-per-minute limit." - }, - "use_openai_responses_path": { - "type": "boolean" - }, - "uses_embed_content": { - "type": "boolean" - }, - "vertex_ai_audio_api": { - "type": "string", - "enum": [ - "lyria_predict", - "lyria_interactions" - ] - }, - "web_search_billing_unit": { - "type": "string", - "description": "Whether web search is billed per query or per prompt.", - "enum": [ - "per_query", - "per_prompt" - ] - } - }, - "additionalProperties": true - } - } -} diff --git a/osv-scanner.toml b/osv-scanner.toml index bb802137f6e..33e2f614edc 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -3,6 +3,11 @@ id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-11-01 reason = "diskcache has no fixed release published; remove this entry once one exists" +[[IgnoredVulns]] +id = "GHSA-v4x8-3wr5-3wq8" +ignoreUntil = 2026-11-20 +reason = "the only fixed filelock is 4.2.1, a new major that requires Python >= 3.11 while litellm still supports 3.10, and the advisory is in ReadWriteLock, which litellm never uses (only FileLock); bump filelock and remove this entry once 3.10 is dropped or a 3.x fix ships" + [[IgnoredVulns]] id = "GHSA-vfj7-8cjw-p6xm" ignoreUntil = 2026-11-03 diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index 2add6868bd9..27681d9116c 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -3,14 +3,14 @@ name = "litellm-core" version = "1.105.0" description = "Library to easily interface with LLM API providers" readme = "README.md" -requires-python = ">=3.10, <3.15" +requires-python = ">=3.10, <3.16" license = "MIT" license-files = ["LICENSE"] authors = [ { name = "BerriAI" }, ] dependencies = [ - "fastuuid>=0.14.0,<1.0", + "fastuuid>=0.14.0,<1.0; python_version < '3.15'", "filelock>=3.16.1,<4.0", "httpx[http2]>=0.28.0,<1.0", "openai>=2.20.0,<4.0.0", @@ -20,16 +20,19 @@ dependencies = [ "packaging>=24.0", "importlib-metadata>=8.0.0,<9.0", "tiktoken>=0.8.0,<1.0; python_version < '3.14'", - "tiktoken>=0.12.0,<1.0; python_version >= '3.14'", + "tiktoken>=0.12.0,<1.0; python_version == '3.14'", + "tiktoken>=0.14.0,<1.0; python_version >= '3.15'", "click>=8.0.0,<9.0", "jinja2>=3.1.6,<4.0", "aiohttp>=3.14.2,<4.0", "anyio>=4.14.0,<5.0", "async-timeout>=4.0.3,<6.0; python_version < '3.11'", "pydantic>=2.11.0,<3.0.0; python_version < '3.14'", - "pydantic>=2.12.0,<3.0.0; python_version >= '3.14'", + "pydantic>=2.12.0,<3.0.0; python_version == '3.14'", + "pydantic>=2.14.0,<3.0.0; python_version >= '3.15'", "pydantic-settings>=2.14.1,<3.0", "jsonschema>=4.0.0,<5.0", + "certifi>=2026.10.19", "typing-extensions>=4.13.0,<5.0", ] diff --git a/pyproject.toml b/pyproject.toml index 9076760153a..fc860a8fdf4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "litellm" version = "1.108.0" description = "Library to easily interface with LLM API providers" readme = "README.md" -requires-python = ">=3.10, <3.15" +requires-python = ">=3.10, <3.16" license = "MIT" license-files = ["LICENSE"] authors = [ @@ -14,7 +14,7 @@ dependencies = [ # deps. Reproducibility for our Docker/CI comes from `uv.lock`. # When changing a floor, verify it installs + imports on every supported # Python with: `uv pip install --resolution=lowest-direct .` - "fastuuid>=0.14.0,<1.0", + "fastuuid>=0.14.0,<1.0; python_version < '3.15'", "filelock>=3.16.1,<4.0", "httpx[http2]>=0.28.0,<1.0", "openai>=2.20.0,<4.0.0", @@ -23,7 +23,8 @@ dependencies = [ "packaging>=24.0", "importlib-metadata>=8.0.0,<9.0", "tiktoken>=0.8.0,<1.0; python_version < '3.14'", - "tiktoken>=0.12.0,<1.0; python_version >= '3.14'", + "tiktoken>=0.12.0,<1.0; python_version == '3.14'", + "tiktoken>=0.14.0,<1.0; python_version >= '3.15'", "tokenizers>=0.21.0,<1.0", "huggingface-hub>=0.34.0,<2.0", "click>=8.0.0,<9.0", @@ -32,10 +33,12 @@ dependencies = [ "anyio>=4.14.0,<5.0", "async-timeout>=4.0.3,<6.0; python_version < '3.11'", "pydantic>=2.11.0,<3.0.0; python_version < '3.14'", - "pydantic>=2.12.0,<3.0.0; python_version >= '3.14'", + "pydantic>=2.12.0,<3.0.0; python_version == '3.14'", + "pydantic>=2.14.0,<3.0.0; python_version >= '3.15'", "pydantic-settings>=2.14.1,<3.0", "jsonschema>=4.0.0,<5.0", - "boto3>=1.43.1,<2.0", + "boto3>=1.43.117,<2.0", + "certifi>=2026.10.19", "typing-extensions>=4.13.0,<5.0", ] @@ -54,17 +57,20 @@ Documentation = "https://docs.litellm.ai" proxy = [ "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", - "granian>=2.7.4,<3.0", + "granian>=2.7.4,<3.0; python_version < '3.15'", + "granian>=2.8.0,<3.0; python_version >= '3.15'", "uvloop>=0.22.1,<1.0; sys_platform != 'win32'", "fastapi>=0.136.3,<1.0", "starlette>=1.0.1,<2.0", "backoff>=2.2.1,<3.0", "pyyaml>=6.0.3,<7.0", "rq>=2.7.0,<3.0", - "orjson>=3.11.6,<4.0", + "orjson>=3.11.6,<4.0; python_version < '3.15'", + "orjson>=3.12.0,<4.0; python_version >= '3.15'", # redis-py's C response parser. It arrives with redis (via rq) either way; naming # it here is what makes redis-py select _HiredisParser instead of the Python one. - "hiredis>=3.0.0,<4.0", + "hiredis>=3.0.0,<4.0; python_version < '3.15'", + "hiredis>=3.5.0,<4.0; python_version >= '3.15'", "apscheduler>=3.11.2,<4.0", "fastapi-sso>=0.19.0,<1.0", "PyJWT>=2.13.0,<3.0", @@ -72,7 +78,7 @@ proxy = [ "cryptography>=49.0.0,<51.0", "pynacl>=1.6.2,<2.0", "websockets>=15.0.1,<16.0", - "boto3>=1.43.1,<2.0", + "boto3>=1.43.117,<2.0", "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=2.3.0,<3", @@ -136,10 +142,10 @@ semantic-router = [ "semantic-router>=0.1.15,<1.0; python_version < '3.14'", "aurelio-sdk>=0.0.19,<1.0; python_version < '3.14'", ] -mlflow = ["mlflow>=3.11.1,<4.0"] +mlflow = ["mlflow>=3.17.2,<4.0"] grpc = [ - # Newest non-yanked release older than the 30-day cutoff. - "grpcio==1.78.0", + # First release whose bundled roots.pem drops the compromised eMudhra roots, ahead of the 30-day cutoff + "grpcio==1.84.1", ] stt-vertex-chirp = [ # Google Cloud Speech-to-Text v2 streaming (gRPC) for Chirp models on @@ -173,19 +179,19 @@ proxy-runtime = [ "google-cloud-speech>=2.40.0,<3.0", "google-genai>=1.37.0,<2.0", "anthropic[vertex]>=0.84.0,<1.0", - "grpcio==1.78.0", + "grpcio==1.84.1", "prometheus-client>=0.20.0,<1.0", "langfuse>=4.7,<5.0", - "opentelemetry-api==1.33.1", - "opentelemetry-sdk==1.33.1", - "opentelemetry-exporter-otlp==1.33.1", - "opentelemetry-instrumentation-fastapi==0.54b1", + "opentelemetry-api==1.45.1", + "opentelemetry-sdk==1.45.1", + "opentelemetry-exporter-otlp==1.45.1", + "opentelemetry-instrumentation-fastapi==0.66b1", "ddtrace>=4.8.2,<5.0", "sentry-sdk>=2.21.0,<3.0", "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", "azure-storage-file-datalake>=12.20.0,<13.0", - "pypdf>=6.16.1,<7.0", + "pypdf>=6.21.2,<7.0", "llm-sandbox>=0.3.39,<1.0", "detect-secrets>=1.5.0,<2.0", ] @@ -209,8 +215,8 @@ dev = [ "tomli==2.4.1; python_version < '3.11'", "pytest-mock==3.15.1", "pytest-asyncio==1.3.0", - "psycopg==3.3.3", - "psycopg-binary==3.3.3", + "psycopg==3.3.6", + "psycopg-binary==3.3.6", "pytest-xdist==3.8.0", "requests-mock==1.12.1", "responses==0.26.0", @@ -222,10 +228,10 @@ dev = [ "types-PyYAML==6.0.12.20250915", "botocore-stubs==1.43.14", "types-boto3[bedrock,bedrock-agent,bedrock-runtime,kms,s3,sagemaker-runtime,sts]==1.43.30", - "opentelemetry-api==1.33.1", - "opentelemetry-sdk==1.33.1", - "opentelemetry-exporter-otlp==1.33.1", - "opentelemetry-instrumentation-fastapi==0.54b1", + "opentelemetry-api==1.45.1", + "opentelemetry-sdk==1.45.1", + "opentelemetry-exporter-otlp==1.45.1", + "opentelemetry-instrumentation-fastapi==0.66b1", "langfuse>=4.7,<5.0", "fastapi-offline==1.7.6", "fakeredis==2.34.1", @@ -254,10 +260,10 @@ proxy-dev = [ "hypercorn==0.17.3", "prometheus-client==0.20.0", "sentry-sdk==2.21.0", - "opentelemetry-api==1.33.1", - "opentelemetry-sdk==1.33.1", - "opentelemetry-exporter-otlp==1.33.1", - "opentelemetry-instrumentation-fastapi==0.54b1", + "opentelemetry-api==1.45.1", + "opentelemetry-sdk==1.45.1", + "opentelemetry-exporter-otlp==1.45.1", + "opentelemetry-instrumentation-fastapi==0.66b1", "azure-identity==1.25.2", "a2a-sdk==1.1.0", "backports-zstd==1.8.0; python_version < '3.14'", @@ -279,7 +285,7 @@ ci = [ "langchain==1.3.9", "lunary==1.4.36; python_version == '3.10'", "lunary==1.4.37; python_version >= '3.11'", - "logfire==4.6.0", + "logfire==5.2.0", "traceloop-sdk==0.34.0", "detect-secrets==1.5.0", "PyGithub==2.8.1", @@ -358,7 +364,7 @@ constraint-dependencies = [ "packaging>=24.0", "soupsieve>=2.8.4", "httplib2>=0.32.0", - "setuptools>=83.0.0", + "setuptools>=84.1.0", ] override-dependencies = [ # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0. @@ -370,6 +376,10 @@ default-groups = ["dev"] required-version = ">=0.10.9" exclude-newer = "3 days" +[tool.uv.exclude-newer-package] +# TODO: remove after 2026-10-25, once brotli 1.2.1 (CVE-2026-111497) clears the 3-day window +brotli = "2026-10-22T14:01:00Z" + [tool.uv.sources] litellm-proxy-extras = { workspace = true } litellm-enterprise = { workspace = true } diff --git a/scripts/sync_together_ai_models.py b/scripts/sync_together_ai_models.py deleted file mode 100644 index 12b128890f1..00000000000 --- a/scripts/sync_together_ai_models.py +++ /dev/null @@ -1,547 +0,0 @@ -"""Sync the together_ai entries of model_prices_and_context_window.json with Together's live serverless catalog. - -Pulls ``GET https://api.together.ai/v1/models?serverless`` plus the deprecations doc, maps API fields onto -registry fields, merges the reviewed capability rules below for everything the API cannot express, and diffs -the result against the registry. Dry run (the default) prints the diff summary and the generated PR body; -``--write`` applies the changes to the root cost map and its ``litellm/`` backup copy. - -Policy highlights: -- Prices arrive per 1M tokens with float artifacts and are normalized to clean per-token values. -- A registry entry absent from the serverless catalog is marked with ``deprecation_date`` from the docs - deprecation table, never deleted; absences with no docs date are surfaced for a human call. -- Availability comes from the API: a model the docs list as removed but the API still serves stays live, - with the conflict surfaced as a warning. -- Manually curated values the API cannot express (``metadata.successor``, ``max_output_tokens`` on existing - entries, capability flags no rule covers) are never overwritten; conflicts are surfaced instead. -""" - -import argparse -import json -import os -import re -import sys -from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field -from pathlib import Path -from types import MappingProxyType -from typing import Final - -import httpx -from pydantic import BaseModel, TypeAdapter, ValidationError - -MODELS_URL: Final = "https://api.together.ai/v1/models?serverless" -DEPRECATIONS_URL: Final = "https://docs.together.ai/docs/deprecations.md" -PROVIDER: Final = "together_ai" -PREFIX: Final = "together_ai/" -SOURCE_URL: Final = "https://docs.together.ai/docs/serverless-models" -COST_MAP_RELPATHS: Final = ( - "model_prices_and_context_window.json", - "litellm/model_prices_and_context_window_backup.json", -) - -TYPE_TO_MODE: Final = MappingProxyType({"chat": "chat", "embedding": "embedding", "moderation": "chat"}) - - -class SyncError(RuntimeError): - pass - - -class CatalogPricing(BaseModel): - input: float - output: float - cached_input: float | None = None - - -class CatalogModel(BaseModel): - id: str - type: str - context_length: int | None = None - pricing: CatalogPricing - - -CATALOG_ADAPTER: Final = TypeAdapter(list[CatalogModel]) - -RegistryEntry = dict[str, object] -CostMap = dict[str, object] - - -@dataclass(frozen=True, slots=True) -class CapabilityRule: - model_id: str - fields: Mapping[str, bool | int] - provenance: str - - -def _rule(model_id: str, provenance: str, **fields: bool | int) -> CapabilityRule: - return CapabilityRule(model_id=model_id, fields=MappingProxyType(dict(fields)), provenance=provenance) - - -_TOOLS: Final = MappingProxyType( - { - "supports_function_calling": True, - "supports_parallel_function_calling": True, - "supports_response_schema": True, - "supports_tool_choice": True, - } -) - -CAPABILITY_RULES: Final = ( - _rule( - "MiniMaxAI/MiniMax-M3", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/minimax-m3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule("Prism-ML/Ternary-Bonsai-27B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "Qwen/Qwen3.5-9B", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/qwen3-5-9b", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule( - "Qwen/Qwen3.6-Plus", - "reviewed for the LIT-5968 backfill; hybrid reasoning model without a documented tools contract", - supports_reasoning=True, - ), - _rule("Qwen/Qwen3.7-Max", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("Qwen/Qwen3.7-Plus", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("Qwen/Qwen3.8-2.4T-A95B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("arize-ai/qwen-2-1.5b-instruct", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "deepseek-ai/DeepSeek-V4-Flash-0731", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-flash", - **_TOOLS, - ), - _rule( - "deepseek-ai/DeepSeek-V4-Pro", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "deepseek-ai/DeepSeek-V4-Pro-0813", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro", - **_TOOLS, - ), - _rule("google/gemma-3n-E4B-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "google/gemma-4-31B-it", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gemma-4-31b-it", - **_TOOLS, - supports_vision=True, - ), - _rule( - "intfloat/multilingual-e5-large-instruct", - "embedding dims per https://huggingface.co/intfloat/multilingual-e5-large-instruct", - output_vector_size=1024, - ), - _rule( - "meta-llama/Llama-3.3-70B-Instruct-Turbo", - "reviewed for the LIT-5968 backfill against https://docs.together.ai/docs/function-calling", - **_TOOLS, - ), - _rule( - "meta-llama/Llama-Guard-4-12B", - "moderation classifier with a chat-shaped API; no tools per the LIT-5968 backfill review", - ), - _rule("meta-models/Muse-Glimmer-30B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "moonshotai/Kimi-K2.7-Code", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k2-7-code", - **_TOOLS, - supports_vision=True, - ), - _rule( - "moonshotai/Kimi-K3", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule( - "nvidia/nemotron-3-ultra-550b-a55b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/nemotron-3-ultra", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "openai/gpt-oss-120b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-120b", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "openai/gpt-oss-20b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-20b", - **_TOOLS, - ), - _rule("pearl-ai/gemma-4-31b-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "thinkingmachines/Inkling", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/inkling", - **_TOOLS, - ), - _rule("thinkingmachines/Inkling-Small", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "zai-org/GLM-5.2", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/glm-5-2;" - " 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.2", - **_TOOLS, - supports_reasoning=True, - max_output_tokens=128000, - max_tokens=128000, - ), - _rule( - "zai-org/GLM-5.3-Flash", - "reviewed for LIT-6489 against https://www.together.ai/models/glm-5-3-flash;" - " 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - max_output_tokens=128000, - max_tokens=128000, - ), -) - -RULES_BY_ID: Final = MappingProxyType({rule.model_id: rule for rule in CAPABILITY_RULES}) - - -@dataclass(frozen=True, slots=True) -class DeprecationDoc: - removal_dates: Mapping[str, str] - redirects: Mapping[str, str] - - -_REDIRECT_ROW: Final = re.compile(r"^\|\s*`([^`]+)`\s*\|\s*`([^`]+)`\s*\|") -_REMOVAL_ROW: Final = re.compile(r"^\|\s*(\d{4}-\d{2}-\d{2})\s*\|\s*`([^`]+)`\s*\|") - - -def _section(markdown: str, heading: str) -> str: - level: Final = heading.split(" ", 1)[0] - start: Final = markdown.find(f"\n{heading}\n") - if start < 0: - return "" - body: Final = markdown[start + 1 + len(heading) :] - next_heading: Final = re.search(rf"^{re.escape(level)} ", body, flags=re.MULTILINE) - return body[: next_heading.start()] if next_heading else body - - -def parse_deprecations(markdown: str) -> DeprecationDoc: - redirect_rows: Final = tuple( - m.groups() - for m in (_REDIRECT_ROW.match(line) for line in _section(markdown, "## Active model redirects").splitlines()) - if m - ) - inference: Final = _section(_section(markdown, "## Deprecation history"), "### Inference") - removal_rows: Final = tuple(m.groups() for m in (_REMOVAL_ROW.match(line) for line in inference.splitlines()) if m) - if not redirect_rows or not removal_rows: - raise SyncError( - "deprecations doc parsed to zero redirect or removal rows; the table format at " - f"{DEPRECATIONS_URL} changed and the parser needs updating" - ) - removal_dates: Final = {model: date for date, model in reversed(removal_rows)} - return DeprecationDoc( - removal_dates=MappingProxyType(dict(reversed(removal_dates.items()))), - redirects=MappingProxyType({original: target for original, target in redirect_rows}), - ) - - -def per_token(price_per_million: float) -> float: - return float(f"{price_per_million / 1e6:.6g}") - - -def _resolve_name(name: str, universe: frozenset[str]) -> str | None: - if name in universe: - return name - suffix_matches: Final = tuple(candidate for candidate in universe if candidate.endswith(f"/{name}")) - return suffix_matches[0] if len(suffix_matches) == 1 else None - - -def resolve_successor(model_id: str, doc: DeprecationDoc, live_ids: frozenset[str]) -> str | None: - canonical: Final = live_ids | frozenset(doc.removal_dates) - redirects: Final = { - (_resolve_name(raw_source, canonical) or raw_source): (_resolve_name(raw_target, canonical) or raw_target) - for raw_source, raw_target in doc.redirects.items() - } - seen: Final = set() - current = model_id # rebind-ok: walks the redirect chain - while current in redirects and current not in seen: - seen.add(current) - current = redirects[current] # rebind-ok: walks the redirect chain - return current if current != model_id and current in live_ids else None - - -@dataclass(frozen=True, slots=True) -class SyncOutcome: - cost_map: CostMap - added: tuple[str, ...] = () - updated: tuple[str, ...] = () - deprecated: tuple[str, ...] = () - reappeared: tuple[str, ...] = () - warnings: tuple[str, ...] = () - skipped_types: Mapping[str, int] = field(default_factory=dict) - - @property - def has_changes(self) -> bool: - return bool(self.added or self.updated or self.deprecated or self.reappeared) - - -def _api_fields(model: CatalogModel) -> RegistryEntry: - cached: Final = model.pricing.cached_input - return { - "input_cost_per_token": per_token(model.pricing.input), - "output_cost_per_token": per_token(model.pricing.output), - **({"cache_read_input_token_cost": per_token(cached), "supports_prompt_caching": True} if cached else {}), - **({"max_input_tokens": model.context_length} if model.context_length is not None else {}), - } - - -def _new_entry(model: CatalogModel, mode: str) -> RegistryEntry: - rule: Final = RULES_BY_ID.get(model.id) - legacy_ceiling: Final = {} if model.context_length is None else {"max_tokens": model.context_length} - merged: Final = { - **_api_fields(model), - **legacy_ceiling, - "litellm_provider": PROVIDER, - "mode": mode, - "source": SOURCE_URL, - **(dict(rule.fields) if rule else {}), - } - return dict(sorted(merged.items())) - - -def _updated_entry(entry: RegistryEntry, model: CatalogModel) -> tuple[RegistryEntry, tuple[str, ...]]: - rule: Final = RULES_BY_ID.get(model.id) - desired: Final = {**_api_fields(model), **(dict(rule.fields) if rule else {})} - dropped: Final = () if model.pricing.cached_input else ("cache_read_input_token_cost", "supports_prompt_caching") - changes: Final = tuple( - f"{name}: {entry.get(name)!r} -> {value!r}" for name, value in desired.items() if entry.get(name) != value - ) + tuple( - f"{name}: {entry[name]!r} removed (no longer in the catalog pricing)" for name in dropped if name in entry - ) - merged: Final = {name: value for name, value in {**entry, **desired}.items() if name not in dropped} - return dict(sorted(merged.items())), changes - - -def _with_new_keys_in_block(original: CostMap, result: CostMap, new_keys: Sequence[str]) -> CostMap: - provider_keys: Final = tuple(key for key in original if key.startswith(PREFIX)) - if not new_keys or not provider_keys: - return result - block_end: Final = provider_keys[-1] - return { - key: value - for existing in original - for key, value in ( - (existing, result[existing]), - *((new, result[new]) for new in sorted(new_keys) if existing == block_end), - ) - } - - -def compute_sync(cost_map: CostMap, catalog: Sequence[CatalogModel], doc: DeprecationDoc) -> SyncOutcome: - live_ids: Final = frozenset(model.id for model in catalog) - token_models: Final = {model.id: model for model in catalog if model.type in TYPE_TO_MODE} - skipped: Final = { - model.type: sum(1 for m in catalog if m.type == model.type) - for model in catalog - if model.type not in TYPE_TO_MODE - } - registry_ids: Final = {key.removeprefix(PREFIX): key for key in cost_map if key.startswith(PREFIX)} - - added: Final[list[str]] = [] - updated: Final[list[str]] = [] - deprecated: Final[list[str]] = [] - reappeared: Final[list[str]] = [] - warnings: Final[list[str]] = [] - result: Final[CostMap] = dict(cost_map) - - for model_id, model in sorted(token_models.items()): - mode: Final = TYPE_TO_MODE[model.type] - key: Final = f"{PREFIX}{model_id}" - if model_id in doc.removal_dates: - warnings.append( - f"`{key}` is listed as removed on {doc.removal_dates[model_id]} in the docs but the serverless " - "catalog still serves it; availability kept from the API" - ) - entry = result.get(key) - if not isinstance(entry, dict): - result[key] = _new_entry(model, mode) - added.append(key) - if model.type == "chat" and model_id not in RULES_BY_ID: - warnings.append( - f"`{key}` added without a capability rule; review its tools/vision/reasoning support and add one" - ) - continue - if entry.get("mode") != mode: - warnings.append( - f"`{key}` has curated mode {entry.get('mode')!r} but the catalog maps to {mode!r}; left unchanged" - ) - new_entry, changes = _updated_entry(entry, model) - if "deprecation_date" in new_entry: - new_entry.pop("deprecation_date") - reappeared.append(key) - if changes: - updated.append(f"{key}: " + "; ".join(changes)) - if changes or key in reappeared: - result[key] = new_entry - - for model_id, key in sorted(registry_ids.items()): - if model_id in token_models: - continue - entry = result.get(key) - if not isinstance(entry, dict): - continue - removal_date: Final = doc.removal_dates.get(model_id) - successor: Final = resolve_successor(model_id, doc, live_ids) - metadata = entry.get("metadata") - curated_successor: Final = metadata.get("successor") if isinstance(metadata, dict) else None - new_entry = dict(entry) - if removal_date is not None and entry.get("deprecation_date") != removal_date: - if "deprecation_date" in entry: - warnings.append( - f"`{key}` has curated deprecation_date {entry.get('deprecation_date')!r} but the docs list " - f"{removal_date!r}; left unchanged" - ) - else: - new_entry["deprecation_date"] = removal_date - if removal_date is None and "deprecation_date" not in entry: - warnings.append( - f"`{key}` is absent from the serverless catalog with no removal date in the docs; " - "needs a human deprecation call" - ) - if successor is not None: - desired_successor: Final = f"{PREFIX}{successor}" - if curated_successor is None: - new_entry["metadata"] = dict( - sorted({**(metadata if isinstance(metadata, dict) else {}), "successor": desired_successor}.items()) - ) - elif curated_successor != desired_successor: - warnings.append( - f"`{key}` has curated successor {curated_successor!r} but the docs redirects resolve to " - f"{desired_successor!r}; left unchanged" - ) - if new_entry != entry: - result[key] = dict(sorted(new_entry.items())) - deprecated.append(f"{key}: " + ", ".join(sorted(set(new_entry) - set(entry)) or ["updated"])) - - return SyncOutcome( - cost_map=_with_new_keys_in_block(cost_map, result, tuple(added)), - added=tuple(added), - updated=tuple(updated), - deprecated=tuple(deprecated), - reappeared=tuple(reappeared), - warnings=tuple(warnings), - skipped_types=MappingProxyType(skipped), - ) - - -def _section_block(title: str, lines: Sequence[str], backtick: bool) -> str: - bullets: Final = "\n".join(f"- `{line}`" if backtick else f"- {line}" for line in lines) or "- none" - return f"### {title} ({len(lines)})\n{bullets}\n" - - -def render_pr_body(outcome: SyncOutcome) -> str: - skipped: Final = ", ".join(f"{kind} ({count})" for kind, count in sorted(outcome.skipped_types.items())) or "none" - return ( - "Automated daily sync of the together_ai entries in model_prices_and_context_window.json against " - f"`GET {MODELS_URL}` and {DEPRECATIONS_URL} by scripts/sync_together_ai_models.py.\n" - "\n" - f"{_section_block('Added', outcome.added, backtick=True)}" - "\n" - f"{_section_block('Updated', outcome.updated, backtick=True)}" - "\n" - f"{_section_block('Marked deprecated', outcome.deprecated, backtick=True)}" - "\n" - f"{_section_block('Returned to the catalog', outcome.reappeared, backtick=True)}" - "\n" - f"{_section_block('Warnings needing a human call', outcome.warnings, backtick=False)}" - "\n" - f"Catalog model types outside the sync's token-pricing scope, skipped: {skipped}\n" - ) - - -def render_summary(outcome: SyncOutcome) -> str: - return ( - f"added={len(outcome.added)} updated={len(outcome.updated)} deprecated={len(outcome.deprecated)} " - f"reappeared={len(outcome.reappeared)} warnings={len(outcome.warnings)}" - ) - - -def load_catalog(raw: bytes) -> list[CatalogModel]: - parsed: Final = json.loads(raw) - entries: Final = parsed.get("data") if isinstance(parsed, dict) else parsed - try: - catalog: Final = CATALOG_ADAPTER.validate_python(entries) - except ValidationError as error: - raise SyncError(f"the catalog response no longer matches the expected shape: {error}") from error - if not any(model.type in TYPE_TO_MODE for model in catalog): - raise SyncError( - "the catalog response contains no token-priced models; refusing to mark the whole registry deprecated" - ) - return catalog - - -def _fetch(url: str, headers: Mapping[str, str]) -> bytes: - response: Final = httpx.get(url, headers=dict(headers), timeout=30, follow_redirects=True) - if response.status_code != 200: - raise SyncError(f"GET {url} returned {response.status_code}") - return response.content - - -def _serialize(cost_map: CostMap) -> str: - return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" - - -def main(argv: Sequence[str]) -> int: - parser: Final = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--write", action="store_true", help="apply the sync to the cost map files (default: dry run)") - parser.add_argument("--models-json", type=Path, help="recorded catalog response to use instead of the live API") - parser.add_argument( - "--deprecations-md", type=Path, help="recorded deprecations doc to use instead of the live docs" - ) - parser.add_argument("--pr-body-file", type=Path, help="write the generated PR body to this path") - parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parent.parent) - args: Final = parser.parse_args(argv) - - if args.models_json is not None: - catalog_raw: Final = args.models_json.read_bytes() - else: - api_key: Final = os.environ.get("TOGETHER_API_KEY") - if not api_key: - raise SyncError("TOGETHER_API_KEY is not set and --models-json was not given") - catalog_raw = _fetch(MODELS_URL, {"Authorization": f"Bearer {api_key}"}) # rebind-ok: branch-dependent source - catalog: Final = load_catalog(catalog_raw) - markdown: Final = ( - args.deprecations_md.read_text() if args.deprecations_md is not None else _fetch(DEPRECATIONS_URL, {}).decode() - ) - doc: Final = parse_deprecations(markdown) - - cost_map_path: Final = args.repo_root / COST_MAP_RELPATHS[0] - cost_map: Final = json.loads(cost_map_path.read_text()) - outcome: Final = compute_sync(cost_map, catalog, doc) - body: Final = render_pr_body(outcome) - - if args.pr_body_file is not None: - args.pr_body_file.write_text(body) - if args.write and outcome.has_changes: - for relpath in COST_MAP_RELPATHS: - (args.repo_root / relpath).write_text(_serialize(outcome.cost_map)) - print(render_summary(outcome)) - print() - print(body) - if not args.write: - print("dry run: no files were touched") - elif not outcome.has_changes: - print("registry already in sync: no files were touched") - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main(sys.argv[1:])) - except SyncError as error: - print(f"SYNC FAILED: {error}", file=sys.stderr) - raise SystemExit(1) from error diff --git a/tests/audio_tests/gemini_tts_speech.mp3 b/tests/audio_tests/gemini_tts_speech.mp3 new file mode 100644 index 00000000000..71a916ad6c0 Binary files /dev/null and b/tests/audio_tests/gemini_tts_speech.mp3 differ diff --git a/tests/audio_tests/gemini_tts_speech.ogg b/tests/audio_tests/gemini_tts_speech.ogg new file mode 100644 index 00000000000..13397fb0c9d Binary files /dev/null and b/tests/audio_tests/gemini_tts_speech.ogg differ diff --git a/tests/benchmarks/test_responses_compaction_benchmarks.py b/tests/benchmarks/test_responses_compaction_benchmarks.py new file mode 100644 index 00000000000..62a5e95255e --- /dev/null +++ b/tests/benchmarks/test_responses_compaction_benchmarks.py @@ -0,0 +1,35 @@ +import json +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.responses.litellm_completion_transformation.compaction_items import ( + with_answered_turns_after_newest_compaction, +) + +COMPACTION_ITEM: Final[Mapping[str, object]] = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} +LONG_TURN_OF_ASSISTANT_TEXT: Final[tuple[Mapping[str, object], ...]] = tuple( + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": f"Step {index}."}]} + for index in range(5_000) +) +CALL: Final[Mapping[str, object]] = { + "type": "function_call", + "call_id": "call_osaka", + "name": "get_weather", + "arguments": "{}", +} +OUTPUT: Final[Mapping[str, object]] = {"type": "function_call_output", "call_id": "call_osaka", "output": "Sunny."} + + +@pytest.mark.benchmark +def test_with_answered_turns_after_newest_compaction_long_turn_split_from_its_output() -> None: + replayed: Final = with_answered_turns_after_newest_compaction( + [*LONG_TURN_OF_ASSISTANT_TEXT, CALL, COMPACTION_ITEM, OUTPUT] + ) + + assert replayed == [COMPACTION_ITEM, *LONG_TURN_OF_ASSISTANT_TEXT, CALL, OUTPUT] diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index e533907659b..cd951c6cd28 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -176,6 +176,7 @@ langgraph: >=1.0.10 # MIT License langgraph-prebuilt: >=1.0.8 # MIT License - https://github.com/langchain-ai/langgraph/blob/main/LICENSE hypothesis: >=6.165.10 # MPL 2.0 license pytest-rerunfailures: >=15.1 # MPL 2.0 license +certifi: >=2026.10.19 # MPL 2.0 license - https://github.com/certifi/python-certifi/blob/master/LICENSE pytest-recording: >=0.13.4 # MIT license expression: >=5.6.0 # MIT License - https://github.com/cognitedata/Expression/blob/main/LICENSE typing-extensions: >=4.13.0 # PSF-2.0 license - https://github.com/python/typing_extensions/blob/main/LICENSE diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 7610d602fff..a5bb5e0b340 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -42,7 +42,7 @@ IGNORE_FUNCTIONS = [ "contains_encrypted_marker", # max depth set. "_rotate_guardrail_row", # bounded by attempts_left. "_delete_nested_value_custom", # max depth set (bounded by number of path segments). - "filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion. + "_filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion. "__getattr__", # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion. "_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation. "_basic_json_schema_validate", # max depth set. diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 55435c51648..853b10ad97b 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -138,6 +138,9 @@ ignored_function_names = [ "_pre_call_check_requested_chat_params", # Tested through the unsupported-param pre-call check tests in test_router.py "_deployment_supports_param", # Tested through the deployment-capability pre-call check tests in test_router.py "_capability_model_and_provider", # Tested through the authenticating-provider pre-call check test in test_router.py + "_route_capability_pin", # Exercised through Decisions routing tests in test_router.py + "_deployment_serves_route", # Exercised through mixed-group Decisions tests in test_router.py + "_without_deployments_that_cannot_serve_route", # Exercised through Decisions routing tests in test_router.py "_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py @@ -146,6 +149,18 @@ ignored_function_names = [ "async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py "_deployment_as_configured_now", # Read by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py "_admit_turn_through_routing_strategy", # Run by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py + "_callback_selectors", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "bind_to_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py + "serve_through_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py + "_request_adaptive_router_hooks", # Run by serve_through_request_logging, tested through the user_config adaptive test + "_release_global_callbacks", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "_routing_callbacks", # Read by the pre-call checks, tested through the user_config limit tests in test_route_llm_request.py + "_request_policies_not_set_globally", # Read by bind_to_request_logging, tested through the user_config budget test in test_route_llm_request.py + "_optional_loggers", # Read by discard and serve_through_request_logging, tested through the budget limiter discard test + "retire_routing_state", # Run by discard, tested through the budget limiter discard test + "_slack_alerting_loggers", # Run by discard, tested through the Slack alerting discard test in test_router_helper_utils.py + "_adaptive_router_hooks", # Run by discard, tested through the adaptive router discard tests in test_router_helper_utils.py + "_global_callback_candidates", # Run by discard, tested through the survivor discard tests in test_router_helper_utils.py "without_deployments_that_cannot_batch", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py "batch_create_routing", # Tested through the mixed-group acreate_batch tests in test_router.py "is_batch_create_routing", # Tested through the chat-after-batch-create test in test_router.py @@ -156,6 +171,7 @@ ignored_function_names = [ "_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py "_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py "_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py + "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py ] diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index d976a2b5ee3..ef57135a4d2 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests diff --git a/tests/e2e/gateway/perf_ci_config.yml b/tests/e2e/gateway/perf_ci_config.yml index 7ab6637e985..767db313847 100644 --- a/tests/e2e/gateway/perf_ci_config.yml +++ b/tests/e2e/gateway/perf_ci_config.yml @@ -46,6 +46,11 @@ model_list: model: openai/gpt-6-luna api_base: os.environ/E2E_PERF_MOCK_OPENAI_URL api_key: perf-mock-not-checked + - model_name: gpt-6-luna-responses + litellm_params: + model: openai/responses/gpt-6-luna + api_base: os.environ/E2E_PERF_MOCK_OPENAI_URL + api_key: perf-mock-not-checked - model_name: claude-sonnet-5 litellm_params: model: anthropic/claude-sonnet-5 @@ -58,6 +63,13 @@ model_list: aws_region_name: us-east-1 aws_access_key_id: perf-mock-not-checked aws_secret_access_key: perf-mock-not-checked + - model_name: claude-sonnet-5-bedrock-converse + litellm_params: + model: bedrock/converse/anthropic.claude-sonnet-5 + aws_bedrock_runtime_endpoint: os.environ/E2E_PERF_MOCK_BEDROCK_URL + aws_region_name: us-east-1 + aws_access_key_id: perf-mock-not-checked + aws_secret_access_key: perf-mock-not-checked guardrails: - guardrail_name: perf-pre litellm_params: diff --git a/tests/e2e/perf/cases/chat_completions_2rps.py b/tests/e2e/perf/cases/chat_completions_2rps.py new file mode 100644 index 00000000000..9792d99c14f --- /dev/null +++ b/tests/e2e/perf/cases/chat_completions_2rps.py @@ -0,0 +1,15 @@ +"""`/v1/chat/completions` at 2 rps: `gpt-6-luna` is served by the OpenAI mock's chat completions +route, `gpt-6-luna-responses` (`openai/responses/gpt-6-luna`) is bridged to its Responses route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "chat_completions_2rps", + "chat_completions", + direct="gpt-6-luna", + bridge="gpt-6-luna-responses", +) diff --git a/tests/e2e/perf/cases/direct_vs_bridge.py b/tests/e2e/perf/cases/direct_vs_bridge.py new file mode 100644 index 00000000000..13f609c8299 --- /dev/null +++ b/tests/e2e/perf/cases/direct_vs_bridge.py @@ -0,0 +1,31 @@ +"""The mixed case's traffic on one public API, half of it on a deployment the gateway serves on +that provider's own API (direct) and half on a deployment it has to translate to another provider +API (bridge). Same 2 rps, request count, shapes, stream share, callers, guardrails and limits as +`mixed_size_requests_2rps`, so a gap between the two halves and the mixed baseline is the +translation's cost.""" + +from __future__ import annotations + +from dataclasses import replace + +from setup.performance_test_case import Endpoint, PerformanceTestCase, VirtualKey + +from .mixed_size_requests_2rps import CASE as MIXED + + +def _scoped(key: VirtualKey, direct: str) -> VirtualKey: + return replace(key, models=(direct,)) if key.models else key + + +def direct_vs_bridge(name: str, endpoint: Endpoint, *, direct: str, bridge: str) -> PerformanceTestCase: + return replace( + MIXED, + name=name, + teams=( + replace(MIXED.team("team-a"), models=(direct, bridge)), + replace(MIXED.team("team-b"), models=(bridge,)), + ), + virtual_keys=tuple(_scoped(key, direct) for key in MIXED.virtual_keys), + model_groups=(direct, bridge), + traffic=replace(MIXED.traffic, endpoints={endpoint: 1.0}), + ) diff --git a/tests/e2e/perf/cases/event_loop_stalls.py b/tests/e2e/perf/cases/event_loop_stalls.py index 0d67a98257f..8a5654cbf7b 100644 --- a/tests/e2e/perf/cases/event_loop_stalls.py +++ b/tests/e2e/perf/cases/event_loop_stalls.py @@ -8,8 +8,8 @@ new hook that deep-copies the request inline all land here as a new row. Runs on `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 +those ten but did in the earlier before/after legs for 46925 and 46967, or in the first scheduled-lane run, 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.""" @@ -44,6 +44,7 @@ CASE: Final = replace( 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), + KnownStall("TaskHandle._run_coro", max_count=2), ), ), ) diff --git a/tests/e2e/perf/cases/messages_2rps.py b/tests/e2e/perf/cases/messages_2rps.py new file mode 100644 index 00000000000..df5d221d430 --- /dev/null +++ b/tests/e2e/perf/cases/messages_2rps.py @@ -0,0 +1,16 @@ +"""`/v1/messages` at 2 rps: `claude-sonnet-5` is served by the Anthropic mock's Messages route, +`claude-sonnet-5-bedrock-converse` (`bedrock/converse/anthropic.claude-sonnet-5`) is bridged to the +Bedrock mock's Converse route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "messages_2rps", + "messages", + direct="claude-sonnet-5", + bridge="claude-sonnet-5-bedrock-converse", +) diff --git a/tests/e2e/perf/cases/mixed_size_requests_ramp.py b/tests/e2e/perf/cases/mixed_size_requests_ramp.py new file mode 100644 index 00000000000..d7db545c031 --- /dev/null +++ b/tests/e2e/perf/cases/mixed_size_requests_ramp.py @@ -0,0 +1,24 @@ +"""The mixed-size traffic stepped up until the pod stops keeping up: the same callers, guardrails, mocks and body +mix (short chats, 300k to 1M-token contexts, tool calls, half streamed) as `mixed_size_requests_2rps`, 2 rps more +every 90 s from 2 to 14 on the same one 4-worker pod. Where the 2 rps case prices one big body on its own, this one +reports `sustained_rps`, the last step the pod held with big bodies overlapping in every worker, so a handler that +re-walks a 1M-token body on the event loop or a worker that holds several of them in memory at once moves the knee +down and fails the floor.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Final + +from setup.performance_test_case import Limit, Ramp, ramped + +from cases.mixed_size_requests_2rps import CASE as MIXED + +RAMP: Final = Ramp(first_rps=2.0, last_rps=14.0, step_rps=2.0, step_seconds=90.0, held_p99_seconds=5.0) + +CASE: Final = replace( + MIXED, + name="mixed_size_requests_ramp", + traffic=ramped(MIXED.traffic, RAMP, seed=9, concurrency=300), + limits=(Limit("sustained_rps", min=8),), +) diff --git a/tests/e2e/perf/cases/responses_2rps.py b/tests/e2e/perf/cases/responses_2rps.py new file mode 100644 index 00000000000..dd55d536eb6 --- /dev/null +++ b/tests/e2e/perf/cases/responses_2rps.py @@ -0,0 +1,15 @@ +"""`/v1/responses` at 2 rps: `gpt-6-luna` is served by the OpenAI mock's Responses route, +`claude-sonnet-5` is bridged to the Anthropic mock's Messages route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "responses_2rps", + "responses", + direct="gpt-6-luna", + bridge="claude-sonnet-5", +) diff --git a/tests/e2e/perf/cases/small_bodies_ramp.py b/tests/e2e/perf/cases/small_bodies_ramp.py new file mode 100644 index 00000000000..64308168cee --- /dev/null +++ b/tests/e2e/perf/cases/small_bodies_ramp.py @@ -0,0 +1,24 @@ +"""The small-bodies traffic stepped up until the pod stops keeping up: the same callers, guardrails, mocks and +2.5k to 4k-token bodies as `small_bodies_5rps`, 10 rps more every 90 s from 10 to 70 on the same one 4-worker pod. +The 5 rps case prices one request with the pod idle; this one reports `sustained_rps`, the last step the pod +held (gateway added p99 under 5 s, the queue still draining), so a change that moves the knee down, a lock held +across an await, a pool that fills, a flush that stalls the loop while dozens of requests are in flight, fails the +floor instead of hiding under a fixed rate the pod still manages.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Final + +from setup.performance_test_case import Limit, Ramp, ramped + +from cases.small_bodies_5rps import CASE as SMALL + +RAMP: Final = Ramp(first_rps=10.0, last_rps=70.0, step_rps=10.0, step_seconds=90.0, held_p99_seconds=5.0) + +CASE: Final = replace( + SMALL, + name="small_bodies_ramp", + traffic=ramped(SMALL.traffic, RAMP, seed=13, concurrency=400), + limits=(Limit("sustained_rps", min=20),), +) diff --git a/tests/e2e/perf/setup/datastore.py b/tests/e2e/perf/setup/datastore.py index 89739b6d13e..6ce37da3554 100644 --- a/tests/e2e/perf/setup/datastore.py +++ b/tests/e2e/perf/setup/datastore.py @@ -22,6 +22,8 @@ from typing import TYPE_CHECKING, Final, Literal from pydantic import TypeAdapter +from setup.series import Recorder, Series + if TYPE_CHECKING: import redis @@ -182,13 +184,22 @@ def slowest_statement(statements: tuple[StatementCalls, ...] | None) -> Statemen class ConnectionsSampler: """`pg_stat_activity` client backends on the gateway's database once a second on a background thread, minus - this sampler's own connection, so the peak is what the gateway (and the test's own reads) held open.""" + this sampler's own connection, so the peak is what the gateway (and the test's own reads) held open. + `series` keeps every count with its time for a ramped run's per-step peaks.""" def __init__(self, url: str) -> None: self._url: Final = url self._stop: Final = threading.Event() self._thread: Final = threading.Thread(target=self._run, name="postgres-connections", daemon=True) - self.peak: int | None = None + self._counts: Final[Recorder[int]] = Recorder() + + @property + def peak(self) -> int | None: + return max((sample.value for sample in self._counts.series.samples), default=None) + + @property + def series(self) -> Series[int]: + return self._counts.series def start(self) -> None: self._thread.start() @@ -203,12 +214,13 @@ class ConnectionsSampler: with psycopg.connect(self._url, autocommit=True) as conn: while not self._stop.is_set(): (count,) = _COMMITS_ROW.validate_python(conn.execute(_CONNECTIONS_QUERY).fetchone()) - self.peak = count if self.peak is None else max(self.peak, count) + self._counts.record(count) _ = self._stop.wait(_CONNECTIONS_SAMPLE_SECONDS) class NoConnectionsSampler: peak: None = None + series: None = None def start(self) -> None: return None diff --git a/tests/e2e/perf/setup/event_loop.py b/tests/e2e/perf/setup/event_loop.py index dee45c374d9..58144424544 100644 --- a/tests/e2e/perf/setup/event_loop.py +++ b/tests/e2e/perf/setup/event_loop.py @@ -14,10 +14,11 @@ a hold that such a pass accounts for is reported under `garbage collection` inst from __future__ import annotations +import json import os import re from collections.abc import Iterable -from dataclasses import dataclass, replace +from dataclasses import asdict, dataclass, replace from datetime import datetime from itertools import groupby from pathlib import Path @@ -79,11 +80,24 @@ class StallRow: total_seconds: float max_seconds: float allowed_by: KnownStall | None + """Every hold in the row, longest first.""" + holds: tuple[float, ...] = () @property def offending(self) -> bool: return self.allowed_by is None or self.count > self.allowed_by.max_count + @property + def over_allowance(self) -> int: + return self.count if self.allowed_by is None else max(0, self.count - self.allowed_by.max_count) + + @property + def note(self) -> str: + allowed: Final = ( + f"{self.count} of {self.allowed_by.max_count} allowed" if self.allowed_by else "not on the allowlist" + ) + return f"{allowed}, max {self.max_seconds:.2f} s" + def _short(path: str) -> str: if "site-packages/" in path: @@ -160,6 +174,7 @@ def summarize(spec: EventLoopStalls, stalls: Iterable[Stall]) -> tuple[StallRow, 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), + holds=tuple(sorted((stall.seconds for stall in group), reverse=True)), ) for (coroutine, where), group in ( (key, tuple(items)) for key, items in groupby(over, key=lambda s: (s.coroutine, s.where)) @@ -181,6 +196,68 @@ def offenders(rows: Iterable[StallRow]) -> tuple[str, ...]: ) +TOP_HOLDS: Final = 10 + + +@dataclass(frozen=True, slots=True) +class Slice: + """One `_breakdown` row in the shape the other metrics use (label, value, share) plus the allowlist verdict the + dashboard shows under the label and colours by. `share` is None where the rows are not parts of a total.""" + + label: str + value: float + share: float | None + note: str | None + over: bool + + +def _label(row: StallRow) -> str: + return f"{row.coroutine} {row.where}".strip() + + +def _slice(row: StallRow, value: float, total: float) -> Slice: + return Slice(label=_label(row), value=value, share=value / total, note=row.note, over=row.offending) + + +def _breakdown(name: str, slices: Iterable[Slice]) -> tuple[tuple[str, str], ...]: + rows: Final = tuple(asdict(row) for row in slices) + return ((f"perf_event_loop_{name}_breakdown", json.dumps(rows)),) if rows else () + + +def _longest_holds(table: tuple[StallRow, ...]) -> tuple[Slice, ...]: + holds: Final = sorted(((seconds, row) for row in table for seconds in row.holds), key=lambda hold: -hold[0]) + return tuple( + Slice(label=_label(row), value=seconds, share=None, note=None, over=row.offending) + for seconds, row in holds[:TOP_HOLDS] + ) + + +def junit_properties(spec: EventLoopStalls, rows: Iterable[StallRow]) -> tuple[tuple[str, float | str], ...]: + """The stall table as `perf_event_loop_*` properties for the dashboard: every hold over the threshold with the + allowlist's total as its limit, the holds past their allowance (limit 0, the case's one assertion), the longest + and the summed hold, each with a `_breakdown` the dashboard opens on click: holds and seconds per location, the + locations past their cap by how far over they are, and the ten longest single holds with where they landed.""" + table: Final = tuple(rows) + holds: Final = sum(row.count for row in table) + offending: Final = sum(row.over_allowance for row in table) + total_seconds: Final = sum(row.total_seconds for row in table) + return ( + ("perf_event_loop_holds", holds), + ("perf_event_loop_holds_limit", sum(known.max_count for known in spec.allowlist)), + ("perf_event_loop_offending_holds", offending), + ("perf_event_loop_offending_holds_limit", 0), + ("perf_event_loop_hold_seconds_max", max((row.max_seconds for row in table), default=0.0)), + ("perf_event_loop_hold_seconds_total", total_seconds), + *_breakdown("holds", (_slice(row, row.count, holds) for row in table)), + *_breakdown( + "offending_holds", + (_slice(row, row.over_allowance, offending) for row in table if row.over_allowance), + ), + *_breakdown("hold_seconds_total", (_slice(row, row.total_seconds, total_seconds) for row in table)), + *_breakdown("hold_seconds_max", _longest_holds(table)), + ) + + 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)" diff --git a/tests/e2e/perf/setup/locustfile.py b/tests/e2e/perf/setup/locustfile.py index 1c38eae22a1..82a5749becf 100644 --- a/tests/e2e/perf/setup/locustfile.py +++ b/tests/e2e/perf/setup/locustfile.py @@ -4,7 +4,8 @@ time to headers, time to first byte, time to first token (streams), total time, and whether it matched. Runs in its own process (`setup/run.py` starts it), never in the gateway's. Reads `PERF_PLAN_FILE` (JSON list of wire requests), writes `PERF_RECORDS_FILE` (one JSON object per -line), paces to `PERF_RPS` over `PERF_USERS` users, and quits once the last request is recorded. +line), paces to `PERF_RPS` over `PERF_USERS` users (or, when the plan carries `send_at_s`, sends each +request at its scheduled second), and quits once the last request is recorded. """ from __future__ import annotations @@ -35,6 +36,8 @@ _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") _PROGRESS: Final = PlanProgress(len(_PLAN)) +_SCHEDULED: Final = bool(_PLAN) and _PLAN[0].send_at_s is not None +_RUN_STARTED: Final = time.monotonic() _STOP_FILE: Final = os.environ.get("PERF_STOP_FILE") BLOCK_REASON: Final = "perf-test-block" _TOKEN_MAX_AGE_S: Final = 240.0 @@ -211,7 +214,8 @@ def _authorization(user: FastHttpUser, item: WireRequest) -> str: class PlanUser(FastHttpUser): """One simulated client; all users share the plan and together hold the planned rate. Each user starts at 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.""" + users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate. A scheduled + plan (a ramp) is paced by the plan instead: a user takes the next request and sleeps until its second.""" network_timeout = float(os.environ.get("PERF_REQUEST_TIMEOUT", "120")) connection_timeout = 30.0 @@ -219,9 +223,12 @@ class PlanUser(FastHttpUser): def on_start(self) -> None: self._last_run = 0.0 self._last_wait = 0.0 - gevent.sleep(_RNG.uniform(0.0, _INTERVAL_S)) + if not _SCHEDULED: + gevent.sleep(_RNG.uniform(0.0, _INTERVAL_S)) def wait_time(self) -> float: + if _SCHEDULED: + return 0.1 target: Final = _INTERVAL_S * _RNG.uniform(1.0 - _JITTER, 1.0 + _JITTER) now: Final = time.monotonic() run_s: Final = now - self._last_run - self._last_wait if self._last_run else 0.0 @@ -237,6 +244,8 @@ class PlanUser(FastHttpUser): if _PROGRESS.drained: _quit(self) return + if item.send_at_s is not None: + gevent.sleep(max(0.0, _RUN_STARTED + item.send_at_s - time.monotonic())) started_at: Final = time.time() started: Final = time.perf_counter() first_byte_s: float | None = None diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index ccdae6398a9..dccc6e272b3 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -7,14 +7,16 @@ import hashlib import json import statistics from collections import Counter, defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from dataclasses import asdict, dataclass -from itertools import chain +from itertools import chain, groupby +from types import MappingProxyType from typing import Final, Literal from e2e_http import NoBody, get_external, post_external, unwrap from memory_readings import RssCapture, RssReading, WorkerKey from pydantic import BaseModel +from typing_extensions import ReadOnly, TypedDict from setup.datastore import ( DbCounters, @@ -27,7 +29,9 @@ from setup.datastore import ( statement_deltas, ) from setup.mocks.reply_common import Provider -from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase +from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase, Ramp +from setup.ramp import StepReading, render_steps, sustained_rps +from setup.series import Series from setup.telemetry import GatewayTelemetry from setup.wire import WireRecord @@ -122,6 +126,10 @@ class Observation: telemetry_before: GatewayTelemetry | None = None telemetry_after: GatewayTelemetry | None = None peak_connections: int | None = None + ramp: Ramp | None = None + cpu_series: Series[float] | None = None + connection_series: Series[int] | None = None + telemetry_series: Series[GatewayTelemetry] | None = None @property def telemetry(self) -> GatewayTelemetry | None: @@ -239,6 +247,67 @@ class Observation: post: Final = Counter(call.marker for call in self.guardrail_calls if call.input_type == "response") return sum(1 for record in self.served if pre[record.marker] != 1 or post[record.marker] != 1) + def steps(self) -> tuple[StepReading, ...]: + """One reading per step of a ramped run: the gateway's added latency over the step's requests, the rate + at which requests completed inside the step's window, the step's unexpected statuses, and what the pod + was doing in that window (CPU per completed request, Redis latency, Postgres connections).""" + ramp: Final = self.ramp + if ramp is None or not self.records: + return () + started: Final = min(record.started_at for record in self.records) + added: Final = _by_step(ramp, self.latency_pairs("gateway_added_seconds")) + first_token: Final = _by_step(ramp, self.latency_pairs("time_to_first_token_seconds")) + completed: Final = Counter( + int((record.started_at + record.total_s - started) // ramp.step_seconds) for record in self.records + ) + unexpected: Final = Counter( + ramp.step_of(record.seq) for record in self.records if not record.blocked and record.status != 200 + ) + return tuple( + StepReading( + rps=rate, + planned=count, + completed_rps=completed[index] / ramp.step_seconds, + gateway_added_p99=percentile(added.get(index, ()), 99), + unexpected=unexpected[index], + held_p99_seconds=ramp.held_p99_seconds, + gateway_added_p50=percentile(added.get(index, ()), 50), + time_to_first_token_p99=percentile(first_token.get(index, ()), 99), + **self._window_readings(started + index * ramp.step_seconds, ramp.step_seconds, completed[index]), + ) + for index, (rate, count) in enumerate(zip(ramp.rates, ramp.counts, strict=True)) + ) + + def _window_readings(self, start: float, seconds: float, completed: int) -> WindowReadings: + end: Final = start + seconds + cpu: Final = None if self.cpu_series is None else sum(self.cpu_series.between(start, end)) + conns: Final = None if self.connection_series is None else self.connection_series.between(start, end) + before: Final = None if self.telemetry_series is None else self.telemetry_series.last_before(start) + after: Final = None if self.telemetry_series is None else self.telemetry_series.last_before(end) + return WindowReadings( + cpu_seconds_per_request=None if cpu is None or completed == 0 else cpu / completed, + redis_latency_mean_seconds=None + if before is None or after is None + else after.minus(before).redis_latency.mean, + postgres_peak_connections=max(conns, default=None) if conns else None, + ) + + +class WindowReadings(TypedDict): + cpu_seconds_per_request: ReadOnly[float | None] + redis_latency_mean_seconds: ReadOnly[float | None] + postgres_peak_connections: ReadOnly[int | None] + + +def _by_step(ramp: Ramp, pairs: Iterable[tuple[WireRecord, float]]) -> Mapping[int, tuple[float, ...]]: + ordered: Final = sorted(pairs, key=lambda pair: pair[0].seq) + return MappingProxyType( + { + step: tuple(value for _, value in group) + for step, group in groupby(ordered, key=lambda pair: ramp.step_of(pair[0].seq)) + } + ) + def _text_sha1(text: str, *, intervened: bool) -> str: original: Final = text.removesuffix(" [post]") if intervened else text @@ -415,6 +484,8 @@ def observe( if metric == "cpu_seconds_per_request": cpu: Final = observation.gateway_cpu_seconds return None if cpu is None or not observation.records else cpu / len(observation.records) + if metric == "sustained_rps": + return None if observation.ramp is None else sustained_rps(observation.steps()) return worst_worker_growth_mb(observation.rss_before, observation.rss_after) @@ -424,6 +495,7 @@ class Row: percentile: int | None observed: float | None limit: float | None + floor: float | None = None detail: str = "" breakdown: tuple[Contributor, ...] = () @@ -437,11 +509,22 @@ class Row: @property def over(self) -> bool: - return self.limit is not None and self.observed is not None and self.observed > self.limit + """Past the limit: above a max, or below a floor.""" + if self.observed is None: + return False + return (self.limit is not None and self.observed > self.limit) or ( + self.floor is not None and self.observed < self.floor + ) @property def missing(self) -> bool: - return self.limit is not None and self.observed is None + return (self.limit is not None or self.floor is not None) and self.observed is None + + @property + def bound(self) -> str: + if self.limit is not None: + return f"{self.limit:.4f}" + return "" if self.floor is None else f">= {self.floor:.4f}" REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( @@ -476,6 +559,8 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ (limit.metric, limit.percentile): limit for limit in case.limits } keys: list[tuple[Metric, int | None]] = list(REPORTED) + if observation.ramp is not None: + keys.append(("sustained_rps", None)) keys.extend(key for key in limits if key not in keys) rows: list[Row] = [] for metric, rank in keys: @@ -486,20 +571,23 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ percentile=rank, observed=observe(metric, rank, observation, case.background_cap_exempt_tables), limit=limit.max if limit else None, + floor=limit.min if limit else None, detail=_detail(metric, observation, case.background_cap_exempt_tables), breakdown=breakdown(metric, rank, observation, case.background_cap_exempt_tables), ) ) - exceeded: Final = [ - f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}{f' ({row.detail})' if row.detail else ''}" - if row.over and row.observed is not None and row.limit is not None - else f"{row.label}: limit set but nothing observed" - for row in rows - if row.over or row.missing - ] + exceeded: Final = [_breach(row) for row in rows if row.over or row.missing] return rows, exceeded +def _breach(row: Row) -> str: + if row.observed is None: + return f"{row.label}: limit set but nothing observed" + if row.floor is not None: + return f"{row.label}: observed {row.observed:.4f} < floor {row.floor:.4f}" + return f"{row.label}: observed {row.observed:.4f} > limit {row.bound}{f' ({row.detail})' if row.detail else ''}" + + def _detail(metric: Metric, observation: Observation, exempt_tables: frozenset[str]) -> str: if metric == "postgres_max_statement_per_request": busiest: Final = observation.busiest_awaited_statement @@ -516,12 +604,13 @@ def _detail(metric: Metric, observation: Observation, exempt_tables: frozenset[s def _row_properties(row: Row) -> tuple[tuple[str, float | str], ...]: observed: Final = ((row.property_name, row.observed),) if row.observed is not None else () limit: Final = ((f"{row.property_name}_limit", row.limit),) if row.limit is not None else () + floor: Final = ((f"{row.property_name}_floor", row.floor),) if row.floor is not None else () contributors: Final = ( ((f"{row.property_name}_breakdown", json.dumps([asdict(item) for item in row.breakdown])),) if row.breakdown else () ) - return observed + limit + contributors + return observed + limit + floor + contributors def junit_properties(rows: Iterable[Row]) -> tuple[tuple[str, float | str], ...]: @@ -546,9 +635,17 @@ def render(case: PerformanceTestCase, observation: Observation, rows: list[Row]) ] for row in rows: observed = "n/a" if row.observed is None else f"{row.observed:.4f}" - limit = "" if row.limit is None else f"{row.limit:.4f}" - flag = " OVER" if row.over else " MISSING" if row.missing else "" - lines.append(f" {row.label:32} {observed:>12} {limit:>12}{flag}") + flag = ( + " UNDER" + if row.over and row.floor is not None + else " OVER" + if row.over + else " MISSING" + if row.missing + else "" + ) + lines.append(f" {row.label:32} {observed:>12} {row.bound:>12}{flag}") + lines.extend(render_steps(observation.steps())) lines.extend(_diagnostics(observation)) lines.extend(_datastore_lines(observation)) for record in failures: @@ -565,6 +662,7 @@ def _diagnostics(observation: Observation) -> list[str]: return [ " diagnostics (not limits):", f" bodies off reply {observation.bodies_off_reply()}", + f" mock routes {_mock_routes(observation)}", f" guardrail calls per request {len(observation.guardrail_calls) / count:.3f}", f" guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}", f" retried share {len(observation.retried_markers) / count:.4f}", @@ -575,6 +673,12 @@ def _diagnostics(observation: Observation) -> list[str]: ] +def _mock_routes(observation: Observation) -> str: + """Which provider route each request landed on, so a direct-vs-bridge case shows its split.""" + counts: Final = Counter(f"{seen.provider}/{seen.route}" for seen in observation.mock_requests) + return ", ".join(f"{route} {count}" for route, count in sorted(counts.items())) or "n/a" + + def _datastore_lines(observation: Observation) -> list[str]: """What the per-request Postgres and Redis counts are made of, so a red run names the statement or command that grew. `plumbing` (BEGIN, COMMIT, session settings) and `unclassified` count in neither limit.""" diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index 3f39b4c38c8..512ef8b911f 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -9,8 +9,10 @@ so a case is a literal a reviewer can read top to bottom. from __future__ import annotations +import bisect from collections.abc import Mapping -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace +from itertools import accumulate from typing import Final, Literal from setup.event_loop import EventLoopStalls @@ -36,6 +38,7 @@ Metric = Literal[ "postgres_active_sessions", "postgres_slowest_statement_seconds", "postgres_peak_connections", + "sustained_rps", ] @@ -97,12 +100,50 @@ SHORT: Final = RequestShape("short", input_tokens=(5_000, 100_000), tools=True) LONG: Final = RequestShape("long", input_tokens=(300_000, 1_000_000), tools=True) +@dataclass(frozen=True, slots=True) +class Ramp: + """Rate stepped up through the run: `first_rps`, then `step_rps` more every `step_seconds`, up to `last_rps`. + Each step is scored on its own and the run reports `sustained_rps`, the last step the pod held before the + first one it did not: a step holds when its `gateway_added_seconds` p99 stayed under `held_p99_seconds`, it + completed at least `HELD_SHARE` of its rate and no more than `HELD_UNEXPECTED_SHARE` of its requests came back + with an unexpected status (`setup/ramp.py`). Requests are sent on a schedule drawn from the plan's seed, not + from per-user intervals, so the rate changes on the step boundary.""" + + first_rps: float + last_rps: float + step_rps: float + step_seconds: float + held_p99_seconds: float + + @property + def rates(self) -> tuple[float, ...]: + count: Final = round((self.last_rps - self.first_rps) / self.step_rps) + 1 + return tuple(self.first_rps + index * self.step_rps for index in range(count)) + + @property + def counts(self) -> tuple[int, ...]: + return tuple(round(rate * self.step_seconds) for rate in self.rates) + + @property + def total_requests(self) -> int: + return sum(self.counts) + + @property + def seconds(self) -> float: + return len(self.rates) * self.step_seconds + + def step_of(self, seq: int) -> int: + """Zero-based step of the request with 1-based plan number `seq`.""" + return bisect.bisect_left(tuple(accumulate(self.counts)), seq) + + @dataclass(frozen=True, slots=True) class Traffic: """The plan: `total_requests` drawn from `seed`, sent at `rps` by `concurrency` simulated - users. Shares are fractions of the requests: `endpoints` and `shapes` must each sum to 1, - `intervene_share` picks non-streaming requests the guardrail rewrites, `block_share` picks - requests it blocks.""" + users, or stepped up through a `ramp` (then `rps` is its first step and `total_requests` its + total, see `ramped`). Shares are fractions of the requests: `endpoints` and `shapes` must each + sum to 1, `intervene_share` picks non-streaming requests the guardrail rewrites, `block_share` + picks requests it blocks.""" seed: int rps: float @@ -113,20 +154,33 @@ class Traffic: intervene_share: float = 0.0 block_share: float = 0.0 concurrency: int = 20 + ramp: Ramp | None = None @property def expected_seconds(self) -> float: - return self.total_requests / self.rps + return self.total_requests / self.rps if self.ramp is None else self.ramp.seconds + + +def ramped(base: Traffic, ramp: Ramp, *, seed: int, concurrency: int) -> Traffic: + """`base`'s request mix sent along `ramp` instead of at one rate. `concurrency` caps the requests in + flight, so size it past the deepest step the ramp is expected to reach.""" + return replace( + base, seed=seed, rps=ramp.first_rps, total_requests=ramp.total_requests, concurrency=concurrency, ramp=ramp + ) @dataclass(frozen=True, slots=True) class Limit: - """One pass/fail rule: `metric` (at `percentile`, or the whole-run value) must not exceed `max`. - Metrics without a Limit are still reported.""" + """One pass/fail rule: `metric` (at `percentile`, or the whole-run value) must not exceed `max`, or must + reach `min` for a floor such as `sustained_rps`. Metrics without a Limit are still reported.""" metric: Metric - max: float + max: float | None = None percentile: int | None = None + min: float | None = None + + def __post_init__(self) -> None: + assert (self.max is None) != (self.min is None), f"{self.metric}: a Limit is either a max or a min" @property def label(self) -> str: @@ -201,6 +255,14 @@ class PerformanceTestCase: assert 0 <= self.traffic.stream_share <= 1, f"{self.name}: stream_share out of range" assert self.traffic.intervene_share + self.traffic.block_share <= 1, f"{self.name}: guardrail shares exceed 1" assert self.traffic.total_requests > 0 and self.traffic.rps > 0, f"{self.name}: empty traffic" + ramp: Final = self.traffic.ramp + assert ramp is None or (ramp.first_rps, ramp.total_requests) == ( + self.traffic.rps, + self.traffic.total_requests, + ), f"{self.name}: traffic rps and total_requests must be the ramp's, build it with ramped()" + assert ramp is None or (0 < ramp.first_rps <= ramp.last_rps and ramp.step_rps > 0 and ramp.step_seconds > 0), ( + f"{self.name}: ramp {ramp} does not step up" + ) assert self.model_groups, f"{self.name}: no model groups" assert 0 <= self.mock_latency.jitter_seconds <= self.mock_latency.seconds or self.mock_latency.seconds == 0, ( f"{self.name}: mock latency jitter {self.mock_latency.jitter_seconds} s exceeds {self.mock_latency.seconds} s" diff --git a/tests/e2e/perf/setup/plan.py b/tests/e2e/perf/setup/plan.py index 6b9a5a19662..4f28736a8e9 100644 --- a/tests/e2e/perf/setup/plan.py +++ b/tests/e2e/perf/setup/plan.py @@ -12,6 +12,7 @@ from dataclasses import dataclass from typing import Final, Literal, TypeVar from setup.performance_test_case import Endpoint, PerformanceTestCase, RequestShape, Team +from setup.ramp import send_offsets Flag = Literal["n", "i", "b"] CallerKind = Literal["jwt", "key"] @@ -39,6 +40,7 @@ class PlannedRequest: flag: Flag shape: RequestShape input_tokens: int + send_at_s: float | None = None @property def blocked(self) -> bool: @@ -98,6 +100,7 @@ def build_plan(case: PerformanceTestCase) -> tuple[PlannedRequest, ...]: case.validate() rng: Final = random.Random(case.traffic.seed) callers: Final = _callers_by_group(case) + offsets: Final = None if case.traffic.ramp is None else send_offsets(case.traffic.ramp, case.traffic.seed) planned: list[PlannedRequest] = [] for seq in range(1, case.traffic.total_requests + 1): endpoint = _pick(rng, case.traffic.endpoints) @@ -117,6 +120,7 @@ def build_plan(case: PerformanceTestCase) -> tuple[PlannedRequest, ...]: flag=flag, shape=shape, input_tokens=rng.randint(low, high), + send_at_s=None if offsets is None else offsets[seq - 1], ) ) return tuple(planned) diff --git a/tests/e2e/perf/setup/podcpu.py b/tests/e2e/perf/setup/podcpu.py index 400e5b7d38b..865fa14efe0 100644 --- a/tests/e2e/perf/setup/podcpu.py +++ b/tests/e2e/perf/setup/podcpu.py @@ -19,6 +19,8 @@ from urllib.parse import quote from e2e_http import AuthHeaders, Success, get_external from pydantic import BaseModel +from setup.series import Recorder, Series + SELECTOR_ENV: Final = "E2E_PERF_GATEWAY_POD_SELECTOR" _SERVICE_ACCOUNT: Final = Path("/var/run/secrets/kubernetes.io/serviceaccount") _INTERVAL_S: Final = 5.0 @@ -49,7 +51,8 @@ def cores(quantity: str) -> float: class GatewayCpuSampler: - """Integrates the matching pods' core use over the time it is running; `seconds` is CPU seconds.""" + """Integrates the matching pods' core use over the time it is running; `seconds` is CPU seconds and + `series` the CPU seconds of each sampling interval, stamped when the interval ended.""" def __init__(self, selector: str, namespace: str, token: str, ca_bundle: str) -> None: self._url: Final = ( @@ -60,12 +63,15 @@ class GatewayCpuSampler: self._ca_bundle: Final = ca_bundle self._stop: Final = threading.Event() self._thread: Final = threading.Thread(target=self._run, name="gateway-cpu-sampler", daemon=True) - self._cpu_seconds = 0.0 - self._samples = 0 + self._intervals: Final[Recorder[float]] = Recorder() @property def seconds(self) -> float | None: - return self._cpu_seconds if self._samples > 0 else None + return sum(sample.value for sample in self._intervals.series.samples) if len(self._intervals) else None + + @property + def series(self) -> Series[float]: + return self._intervals.series def _read_cores(self) -> float | None: result: Final = get_external( @@ -81,8 +87,7 @@ class GatewayCpuSampler: now: float = time.monotonic() current: float | None = self._read_cores() if current is not None: - self._cpu_seconds += current * (now - last) - self._samples += 1 + self._intervals.record(current * (now - last)) last = now # rebind-ok: the sampling clock advances every tick def start(self) -> None: @@ -95,6 +100,7 @@ class GatewayCpuSampler: class NoSampler: seconds: Final[None] = None + series: Final[None] = None def start(self) -> None: return None diff --git a/tests/e2e/perf/setup/ramp.py b/tests/e2e/perf/setup/ramp.py new file mode 100644 index 00000000000..217e140d470 --- /dev/null +++ b/tests/e2e/perf/setup/ramp.py @@ -0,0 +1,116 @@ +"""A ramped run step by step: the send schedule its plan carries (`send_offsets`), one `StepReading` per step +and `sustained_rps`, the last step the pod held before the first it did not. `Ramp` itself sits in +`performance_test_case.py` with the rest of the case schema.""" + +from __future__ import annotations + +import random +from collections.abc import Iterable, Iterator +from dataclasses import dataclass +from itertools import chain, takewhile +from typing import Final + +from setup.performance_test_case import Ramp + +HELD_SHARE: Final = 0.9 +"""A step that completed fewer than this share of its rate is a queue that stopped draining.""" +HELD_UNEXPECTED_SHARE: Final = 0.01 +"""A step with more than this share of unexpected statuses broke, whatever its latency.""" +_JITTER: Final = 0.25 + + +def _step_offsets(ramp: Ramp, index: int, count: int, rng: random.Random) -> Iterator[float]: + for slot in range(count): + yield (index + (slot + 0.5 + rng.uniform(-_JITTER, _JITTER)) / count) * ramp.step_seconds + + +def send_offsets(ramp: Ramp, seed: int) -> tuple[float, ...]: + """Seconds after the run starts to send each request of the plan, in plan order: a step's requests spread + evenly over its window, each moved by up to a quarter of a slot either way, drawn from `seed` so a rerun + sends at the same seconds.""" + rng: Final = random.Random(f"ramp-{seed}") + return tuple(chain.from_iterable(_step_offsets(ramp, index, count, rng) for index, count in enumerate(ramp.counts))) + + +@dataclass(frozen=True, slots=True) +class StepReading: + rps: float + planned: int + completed_rps: float + gateway_added_p99: float | None + unexpected: int + held_p99_seconds: float + gateway_added_p50: float | None = None + time_to_first_token_p99: float | None = None + cpu_seconds_per_request: float | None = None + redis_latency_mean_seconds: float | None = None + postgres_peak_connections: int | None = None + """The hold verdict reads the first six. The rest say what the pod was doing on that step: CPU per request + from the pod sampler, Redis latency from the gateway's `/metrics`, connections from `pg_stat_activity`, + all cut to the step's window, so a latency jump can be told apart as CPU, waiting or neither.""" + + @property + def held(self) -> bool: + return ( + self.gateway_added_p99 is not None + and self.gateway_added_p99 <= self.held_p99_seconds + and self.completed_rps >= HELD_SHARE * self.rps + and self.unexpected <= HELD_UNEXPECTED_SHARE * self.planned + ) + + +def sustained_rps(steps: Iterable[StepReading]) -> float: + """The rate of the last step that held before the first that did not, 0 when the first step did not hold.""" + held: Final = tuple(takewhile(lambda step: step.held, steps)) + return held[-1].rps if held else 0.0 + + +def step_properties(steps: Iterable[StepReading]) -> tuple[tuple[str, float], ...]: + """The per-step readings as JUnit properties (`perf_ramp_step__`, every reading the step has) + plus the hold bound, so the dashboard can draw the run's curves next to the sustained rate.""" + rows: Final = tuple(steps) + hold: Final = (("perf_ramp_hold_p99_seconds", rows[0].held_p99_seconds),) if rows else () + return (*hold, *chain.from_iterable(_step_properties(step) for step in rows)) + + +def _step_properties(step: StepReading) -> Iterator[tuple[str, float]]: + prefix: Final = f"perf_ramp_step_{step.rps:g}" + readings: Final[tuple[tuple[str, float | None], ...]] = ( + ("gateway_added_p99", step.gateway_added_p99), + ("completed_rps", step.completed_rps), + ("unexpected", float(step.unexpected)), + ("gateway_added_p50", step.gateway_added_p50), + ("time_to_first_token_p99", step.time_to_first_token_p99), + ("cpu_seconds_per_request", step.cpu_seconds_per_request), + ("redis_latency_mean_seconds", step.redis_latency_mean_seconds), + ( + "postgres_peak_connections", + None if step.postgres_peak_connections is None else float(step.postgres_peak_connections), + ), + ) + return ((f"{prefix}_{name}", value) for name, value in readings if value is not None) + + +def render_steps(steps: Iterable[StepReading]) -> tuple[str, ...]: + rows: Final = tuple(steps) + if not rows: + return () + header: Final = ( + f" steps: {'rps':>6} {'planned':>8} {'completed rps':>14} {'gateway added p50':>18} {'p99':>8}" + f" {'ttft p99':>9} {'unexpected':>11} {'cpu s/req':>10} {'redis mean':>11} {'pg conns':>9}" + ) + return (header, *(_step_line(step) for step in rows)) + + +def _seconds(value: float | None) -> str: + return "n/a" if value is None else f"{value:.4f}" + + +def _step_line(step: StepReading) -> str: + conns: Final = "n/a" if step.postgres_peak_connections is None else str(step.postgres_peak_connections) + return ( + f" {step.rps:6.1f} {step.planned:8d} {step.completed_rps:14.2f} {_seconds(step.gateway_added_p50):>18}" + f" {_seconds(step.gateway_added_p99):>8} {_seconds(step.time_to_first_token_p99):>9} {step.unexpected:11d}" + f" {_seconds(step.cpu_seconds_per_request):>10} {_seconds(step.redis_latency_mean_seconds):>11} {conns:>9}" + f" {'held' if step.held else 'broke'}" + ) diff --git a/tests/e2e/perf/setup/run.py b/tests/e2e/perf/setup/run.py index 272c64754c4..ec2927e90af 100644 --- a/tests/e2e/perf/setup/run.py +++ b/tests/e2e/perf/setup/run.py @@ -63,6 +63,7 @@ def wire_request(item: PlannedRequest, setup: CaseSetup) -> WireRequest: token_source=setup.token_source(item.caller), stream=item.stream, blocked=item.blocked, + send_at_s=item.send_at_s, ) diff --git a/tests/e2e/perf/setup/series.py b/tests/e2e/perf/setup/series.py new file mode 100644 index 00000000000..78e5cf272e1 --- /dev/null +++ b/tests/e2e/perf/setup/series.py @@ -0,0 +1,48 @@ +"""Readings a background sampler takes while the traffic runs, each stamped with the wall clock +(`time.time()`, the clock the Locust records carry in `started_at`), so a ramped run can cut them into its +steps' windows after the fact.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Final, Generic, TypeVar + +T = TypeVar("T") + + +@dataclass(frozen=True, slots=True) +class Sample(Generic[T]): + at: float + value: T + + +@dataclass(frozen=True, slots=True) +class Series(Generic[T]): + samples: tuple[Sample[T], ...] + + def between(self, start: float, end: float) -> tuple[T, ...]: + """The values stamped inside `[start, end)`, in time order.""" + return tuple(sample.value for sample in self.samples if start <= sample.at < end) + + def last_before(self, at: float) -> T | None: + """The newest value stamped before `at`, None when the series starts later.""" + earlier: Final = tuple(sample.value for sample in self.samples if sample.at < at) + return earlier[-1] if earlier else None + + +class Recorder(Generic[T]): + """What a sampler thread appends to; `series` is a frozen copy for the observation.""" + + def __init__(self) -> None: + self._samples: list[Sample[T]] = [] + + def record(self, value: T) -> None: + self._samples.append(Sample(at=time.time(), value=value)) + + @property + def series(self) -> Series[T]: + return Series(samples=tuple(self._samples)) + + def __len__(self) -> int: + return len(self._samples) diff --git a/tests/e2e/perf/setup/telemetry.py b/tests/e2e/perf/setup/telemetry.py index e52bcabafda..185cc788601 100644 --- a/tests/e2e/perf/setup/telemetry.py +++ b/tests/e2e/perf/setup/telemetry.py @@ -1,5 +1,5 @@ -"""The gateway's own `/metrics`, read before and after a run for the numbers nothing outside the gateway -can see. Today that is `litellm_redis_latency`, the histogram customers' Grafana panels show when Redis +"""The gateway's own `/metrics`, read before and after a run (and every few seconds during a ramped run, see +`telemetry_sampler`) for the numbers nothing outside the gateway can see. Today that is `litellm_redis_latency`, the histogram customers' Grafana panels show when Redis slows the gateway down. The gateway's latency and request counts are read per request from its response headers instead (`setup/wire.py`), which gives percentiles where a histogram delta only gives a mean. @@ -10,7 +10,9 @@ on psutil (`proxy_usage.py`).""" from __future__ import annotations import math -from collections.abc import Iterable +import threading +from collections.abc import Generator, Iterable +from contextlib import contextmanager from dataclasses import dataclass from itertools import chain from typing import Final @@ -20,6 +22,10 @@ from prometheus_client.parser import text_string_to_metric_families from prometheus_client.samples import Sample from proxy_client import ProxyClient +from setup.series import Recorder, Series + +_SAMPLE_SECONDS: Final = 10.0 + @dataclass(frozen=True, slots=True) class HistogramTotals: @@ -63,3 +69,42 @@ def read_gateway_telemetry(proxy: ProxyClient) -> GatewayTelemetry | None: case every limit on a gateway-reported metric fails as missing rather than reading a silent zero.""" scrape: Final = proxy.probe("/metrics", params=NoBody()) return parse_gateway_telemetry(scrape.body) if scrape.status_code == 200 else None + + +class TelemetrySampler: + """`/metrics` every `_SAMPLE_SECONDS` on a background thread while traffic runs, so a ramped run can read + the Redis latency of each step's window as the difference between the scrapes around it.""" + + def __init__(self, proxy: ProxyClient) -> None: + self._proxy: Final = proxy + self._stop: Final = threading.Event() + self._thread: Final = threading.Thread(target=self._run, name="gateway-telemetry", daemon=True) + self._scrapes: Final[Recorder[GatewayTelemetry]] = Recorder() + + @property + def series(self) -> Series[GatewayTelemetry]: + return self._scrapes.series + + def _run(self) -> None: + while not self._stop.is_set(): + scrape: GatewayTelemetry | None = read_gateway_telemetry(self._proxy) + if scrape is not None: + self._scrapes.record(scrape) + _ = self._stop.wait(_SAMPLE_SECONDS) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=30) + + +@contextmanager +def telemetry_sampler(proxy: ProxyClient) -> Generator[TelemetrySampler]: + sampler: Final = TelemetrySampler(proxy) + sampler.start() + try: + yield sampler + finally: + sampler.stop() diff --git a/tests/e2e/perf/setup/upgrade.py b/tests/e2e/perf/setup/upgrade.py index 36efc16e1fe..c9114ccbff2 100644 --- a/tests/e2e/perf/setup/upgrade.py +++ b/tests/e2e/perf/setup/upgrade.py @@ -145,7 +145,9 @@ class UpgradeMonitor: 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.status_limit: Final = next( + limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count" and limit.max is not None + ) self.pending: tuple[WireRecord, ...] = () self.processed = 0 self.successes = 0 diff --git a/tests/e2e/perf/setup/wire.py b/tests/e2e/perf/setup/wire.py index 18351c18c6b..a1833fbd44b 100644 --- a/tests/e2e/perf/setup/wire.py +++ b/tests/e2e/perf/setup/wire.py @@ -51,6 +51,8 @@ class WireRequest(BaseModel): token_source: TokenSource | None = None stream: bool blocked: bool + send_at_s: float | None = None + """Seconds after the run starts to send this request, set when the case ramps; otherwise the users pace.""" class WireRecord(BaseModel): diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index 49b23ddfc57..db657a04b00 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -29,6 +29,7 @@ from setup.event_loop import ( render_stalls, summarize, ) +from setup.event_loop import junit_properties as stall_properties from setup.observe import ( Observation, evaluate, @@ -43,8 +44,9 @@ from setup.performance_test_case import MockLatency, PerformanceTestCase from setup.plan import PlannedRequest, build_plan from setup.podcpu import gateway_cpu_sampler from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url +from setup.ramp import step_properties from setup.run import run_plan -from setup.telemetry import read_gateway_telemetry +from setup.telemetry import read_gateway_telemetry, telemetry_sampler pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] _SPEND_SAMPLE: Final = 40 @@ -114,7 +116,12 @@ def test_perf_case( 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, connections_sampler() as connections: + with ( + gateway_cpu_sampler() as cpu, + redis_monitor() as monitor, + connections_sampler() as connections, + telemetry_sampler(proxy) as telemetry, + ): outcome: Final = run_plan(case, plan, setup) run_finished: Final = time.monotonic() db_after: Final = read_db_counters() @@ -150,6 +157,10 @@ def test_perf_case( telemetry_before=telemetry_before, telemetry_after=telemetry_after, peak_connections=connections.peak, + ramp=case.traffic.ramp, + cpu_series=cpu.series, + connection_series=connections.series, + telemetry_series=telemetry.series, ) rows, exceeded_limits = evaluate(case, observation) stall_rows: Final = _event_loop_rows(case, traffic_started_at) if case.event_loop is not None else () @@ -158,7 +169,10 @@ def test_perf_case( item.user_properties[:] = [ entry for entry in item.user_properties if not entry[0].startswith("perf_") ] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set - item.user_properties.extend(junit_properties(rows)) + item.user_properties.extend( + stall_properties(case.event_loop, stall_rows) if case.event_loop is not None else junit_properties(rows) + ) + item.user_properties.extend(step_properties(observation.steps())) print(render(case, observation, rows)) if case.event_loop is not None: print(render_stalls(case.event_loop, stall_rows)) diff --git a/tests/e2e/ui/globalSetup.ts b/tests/e2e/ui/globalSetup.ts index d97789c227e..e18d1be390c 100644 --- a/tests/e2e/ui/globalSetup.ts +++ b/tests/e2e/ui/globalSetup.ts @@ -2,7 +2,7 @@ import { chromium, expect, request } from "@playwright/test"; import { users, Role, STORAGE_PATHS } from "./fixtures/users"; import { ARTIFACT_DIR, UI_BASE_URL } from "./constants"; import { expectUnrestrictedDashboard, setInvitedUserPassword } from "./helpers/userOnboarding"; -import { hideLiteAdmin } from "./helpers/navigation"; +import { setAriaHidden } from "./helpers/navigation"; import { masterKey as getMasterKey } from "./helpers/traffic"; import * as fs from "fs"; import * as path from "path"; @@ -78,7 +78,7 @@ async function globalSetup() { await dismiss.click(); } if (role === Role.ProxyAdmin) { - await hideLiteAdmin(page); + await setAriaHidden(page, true); } // The login flow stores a post-login return URL in the litellm_return_url // cookie. If the snapshot captures it before the app consumes it, every diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index 78288fab743..be46fcf0395 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -62,18 +62,29 @@ export async function dismissFeedbackPopup(page: PlaywrightPage): Promise } } -export async function hideLiteAdmin(page: PlaywrightPage): Promise { +export async function setAriaHidden(page: PlaywrightPage, hidden: boolean): Promise { await page.getByRole("button", { name: /Account menu/i }).click(); const panel = page.getByTestId("sidebar-account-menu-panel"); await expect(panel).toBeVisible({ timeout: 5_000 }); - const toggle = panel.getByRole("switch", { name: "Toggle hide LiteAdmin" }); - if ((await toggle.getAttribute("aria-checked")) !== "true") { - await toggle.click(); + const premium = await panel.getByText("Premium", { exact: true }).isVisible(); + const toggle = panel.getByRole("switch", { name: "Toggle hide Aria Agent", exact: true }); + if (premium) { + await expect(toggle).toBeVisible(); + if ((await toggle.getAttribute("aria-checked")) !== String(hidden)) { + await toggle.click(); + } + await expect(toggle).toHaveAttribute("aria-checked", String(hidden)); + } else { + await expect(toggle).toHaveCount(0); } - await expect(toggle).toHaveAttribute("aria-checked", "true"); await page.keyboard.press("Escape"); await expect(panel).toBeHidden(); - await expect(page.getByRole("button", { name: "LiteAdmin", exact: true })).toBeHidden(); + const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + if (premium && hidden) { + await expect(aria).toBeHidden(); + } else { + await expect(aria).toBeVisible(); + } } /** diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts index 416e1c3171f..838a829a3e2 100644 --- a/tests/e2e/ui/tests/logs/logsPagination.spec.ts +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -1,8 +1,14 @@ import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; -import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { navigateToPage, dismissFeedbackPopup, setAriaHidden } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; -import { CHAT_MODEL_A, createVirtualKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic"; +import { + CHAT_MODEL_A, + createVirtualKey, + deleteVirtualKey, + sendChatCompletion, + waitForSpendLog, +} from "../../helpers/traffic"; /** * Session-grouped pagination (#38060): a page of N rows must render exactly N session rows, a @@ -45,6 +51,57 @@ async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): await page.getByRole("option", { name: size, exact: true }).click(); } +test.describe("Logs pagination with Aria in the header", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + for (const viewport of [ + { width: 1440, height: 900 }, + { width: 390, height: 844 }, + ]) { + test(`keeps pagination reachable and Aria available at ${viewport.width}px`, async ({ page, request }) => { + await page.setViewportSize(viewport); + const alias = `e2e-aria-header-${uniqueSuffix()}`; + const key = await createVirtualKey(request, { key_alias: alias }); + try { + const ids = await Promise.all( + Array.from({ length: 26 }, (_, index) => + sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `${alias}-${index}`, + apiKey: key.key, + }), + ), + ); + await Promise.all(ids.map((id) => waitForSpendLog(request, id))); + await openLogs(page); + await setAriaHidden(page, false); + await applyKeyAliasFilter(page, await openFilterDrawer(page), alias); + await setRowsPerPage(page, "25"); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + await visibleTestId(page, "pagination-next").click(); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-26 of 26"); + + const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + await aria.click(); + await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); + await page.keyboard.press("Escape"); + await expect(aria).toBeFocused(); + await visibleTestId(page, "pagination-first").click(); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { + await navigateToPage(page, destination); + await aria.click(); + await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); + await page.keyboard.press("Escape"); + await expect(aria).toBeFocused(); + } + } finally { + await deleteVirtualKey(request, key.token); + } + }); + } +}); + test.describe("Logs page session-grouped pagination", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); diff --git a/tests/e2e_harness/perf/setup/test_event_loop.py b/tests/e2e_harness/perf/setup/test_event_loop.py index 57868305716..75989b38275 100644 --- a/tests/e2e_harness/perf/setup/test_event_loop.py +++ b/tests/e2e_harness/perf/setup/test_event_loop.py @@ -1,8 +1,10 @@ from __future__ import annotations +import math from typing import Final -from setup.event_loop import EventLoopStalls, KnownStall, offenders, parse_stalls, summarize +from pydantic import BaseModel, ConfigDict, TypeAdapter +from setup.event_loop import EventLoopStalls, KnownStall, junit_properties, offenders, parse_stalls, summarize _HANDLER: Final = ( "Executing list[_Slice]: + return TypeAdapter(list[_Slice]).validate_json(str(properties[f"perf_event_loop_{name}_breakdown"])) + + 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()" @@ -95,3 +111,72 @@ def test_a_collector_pass_explains_only_the_next_hold_and_only_one_at_least_as_l ("RequestResponseCycle.run_asgi()", False), ] assert summarize(_SPEC, parse_stalls((_GC_PASS,))) == () + + +def test_junit_properties_count_every_hold_and_only_the_holds_past_their_allowance_as_offending() -> None: + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, _STREAM, _HANDLER, _HANDLER, _HANDLER))) + properties: Final = dict(junit_properties(_SPEC, rows)) + assert properties["perf_event_loop_holds"] == 6 + assert properties["perf_event_loop_holds_limit"] == 6 + assert properties["perf_event_loop_offending_holds"] == 3 + assert properties["perf_event_loop_offending_holds_limit"] == 0 + assert properties["perf_event_loop_hold_seconds_max"] == 0.731 + assert math.isclose(float(properties["perf_event_loop_hold_seconds_total"]), 0.302 + 2 * 0.412 + 3 * 0.731) + breakdown: Final = _breakdown(properties, "holds") + assert [(slice.label, slice.value, slice.note, slice.over) for slice in breakdown] == [ + ( + "RequestResponseCycle.run_asgi() uvicorn/protocols/http/h11_impl.py:416", + 3, + "3 of 2 allowed, max 0.73 s", + True, + ), + ("TaskHandle._run_coro() anyio/_core/_tasks.py:275", 2, "not on the allowlist, max 0.41 s", True), + ("LifespanOn.main() uvicorn/lifespan/on.py:86", 1, "1 of 3 allowed, max 0.30 s", False), + ] + assert math.isclose(sum(slice.share or 0.0 for slice in breakdown), 1.0) + + +def test_the_offending_breakdown_lists_only_the_locations_past_their_cap_by_how_far_over_they_are() -> None: + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, _STREAM, _HANDLER, _HANDLER, _HANDLER))) + offending: Final = _breakdown(dict(junit_properties(_SPEC, rows)), "offending_holds") + assert [(slice.label, slice.value, slice.share, slice.over) for slice in offending] == [ + ("RequestResponseCycle.run_asgi() uvicorn/protocols/http/h11_impl.py:416", 1, 1 / 3, True), + ("TaskHandle._run_coro() anyio/_core/_tasks.py:275", 2, 2 / 3, True), + ] + clean: Final = dict(junit_properties(_SPEC, summarize(_SPEC, parse_stalls((_HANDLER, _HANDLER))))) + assert clean["perf_event_loop_offending_holds"] == 0 + assert "perf_event_loop_offending_holds_breakdown" not in clean + + +def test_the_seconds_breakdowns_carry_the_time_per_location_and_the_longest_single_holds() -> None: + long_handler: Final = _HANDLER.replace("0.731", "0.950") + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, long_handler, _HANDLER))) + properties: Final = dict(junit_properties(_SPEC, rows)) + per_location: Final = _breakdown(properties, "hold_seconds_total") + assert [(slice.label.split()[0], slice.value) for slice in per_location] == [ + ("RequestResponseCycle.run_asgi()", 0.950 + 0.731), + ("TaskHandle._run_coro()", 0.412), + ("LifespanOn.main()", 0.302), + ] + assert math.isclose(sum(slice.share or 0.0 for slice in per_location), 1.0) + longest: Final = _breakdown(properties, "hold_seconds_max") + assert [(slice.label.split()[0], slice.value, slice.share, slice.over) for slice in longest] == [ + ("RequestResponseCycle.run_asgi()", 0.950, None, False), + ("RequestResponseCycle.run_asgi()", 0.731, None, False), + ("TaskHandle._run_coro()", 0.412, None, True), + ("LifespanOn.main()", 0.302, None, False), + ] + + +def test_the_longest_holds_breakdown_stops_at_ten() -> None: + rows: Final = summarize(_SPEC, parse_stalls(tuple(_STARTUP for _ in range(12)))) + longest: Final = _breakdown(dict(junit_properties(_SPEC, rows)), "hold_seconds_max") + assert len(longest) == 10 + + +def test_junit_properties_without_holds_report_zeros_and_no_breakdown() -> None: + properties: Final = dict(junit_properties(_SPEC, ())) + assert properties["perf_event_loop_holds"] == 0 + assert properties["perf_event_loop_offending_holds"] == 0 + assert properties["perf_event_loop_hold_seconds_max"] == 0 + assert not [name for name in properties if name.endswith("_breakdown")] diff --git a/tests/e2e_harness/perf/setup/test_ramp.py b/tests/e2e_harness/perf/setup/test_ramp.py new file mode 100644 index 00000000000..6d6c1640b59 --- /dev/null +++ b/tests/e2e_harness/perf/setup/test_ramp.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import math +from dataclasses import replace +from typing import Final + +from cases.small_bodies_5rps import CASE +from memory_readings import RssCapture +from setup.observe import Observation, evaluate, junit_properties +from setup.performance_test_case import Limit, Ramp, ramped +from setup.plan import build_plan +from setup.ramp import StepReading, send_offsets, step_properties, sustained_rps +from setup.series import Sample, Series +from setup.telemetry import GatewayTelemetry, HistogramTotals +from setup.wire import WireRecord + +NO_RSS: Final = RssCapture(readings=(), failures=()) + + +def test_send_offsets_put_each_step_s_requests_inside_its_window_in_plan_order() -> None: + ramp: Final = Ramp(first_rps=10.0, last_rps=30.0, step_rps=10.0, step_seconds=10.0, held_p99_seconds=5.0) + offsets: Final = send_offsets(ramp, seed=13) + assert ramp.counts == (100, 200, 300) + assert len(offsets) == ramp.total_requests + assert offsets == tuple(sorted(offsets)) + for seq, offset in enumerate(offsets, start=1): + step = ramp.step_of(seq) + assert step * 10.0 <= offset < (step + 1) * 10.0, f"request {seq} at {offset:.2f}s is outside step {step}" + assert send_offsets(ramp, seed=13) == offsets + assert send_offsets(ramp, seed=14) != offsets + + +def test_build_plan_carries_the_schedule_only_for_a_ramped_case() -> None: + ramp: Final = Ramp(first_rps=2.0, last_rps=4.0, step_rps=2.0, step_seconds=5.0, held_p99_seconds=5.0) + ramped_case: Final = replace(CASE, traffic=ramped(CASE.traffic, ramp, seed=3, concurrency=20)) + assert ramped_case.traffic.expected_seconds == 10.0 + assert [item.send_at_s for item in build_plan(ramped_case)] == list(send_offsets(ramp, seed=3)) + assert all(item.send_at_s is None for item in build_plan(CASE)) + + +def _step(rps: float, p99: float, completed_rps: float | None = None, unexpected: int = 0) -> StepReading: + return StepReading( + rps=rps, + planned=int(rps * 90), + completed_rps=rps if completed_rps is None else completed_rps, + gateway_added_p99=p99, + unexpected=unexpected, + held_p99_seconds=5.0, + ) + + +def test_sustained_rps_is_the_last_step_held_before_the_first_that_broke() -> None: + assert sustained_rps((_step(10, 0.5), _step(20, 1.0), _step(30, 9.0), _step(40, 0.7))) == 20 + assert sustained_rps((_step(10, 0.5), _step(20, 1.0, completed_rps=17.0), _step(30, 1.0))) == 10 + assert sustained_rps((_step(10, 0.5), _step(20, 1.0, unexpected=19), _step(30, 1.0))) == 10 + assert sustained_rps((_step(10, 0.5), _step(20, 1.0, unexpected=18), _step(30, 1.0))) == 30 + assert sustained_rps((_step(10, 6.0), _step(20, 1.0))) == 0 + + +def _record(seq: int, started_at: float, total_s: float) -> WireRecord: + return WireRecord( + seq=seq, + marker=f"perfreq-{seq:06d}-n", + caller="key-a1", + endpoint="chat_completions", + stream=False, + blocked=False, + started_at=started_at, + status=200, + header_s=total_s, + first_byte_s=total_s, + first_token_s=None, + total_s=total_s, + text="", + ok=True, + mock_delay_s=1.0, + call_id=None, + error=None, + ) + + +def test_a_floor_fails_when_the_second_step_s_latency_breaks_it_and_the_properties_carry_the_floor() -> None: + ramp: Final = Ramp(first_rps=10.0, last_rps=20.0, step_rps=10.0, step_seconds=100.0, held_p99_seconds=2.0) + case: Final = replace( + CASE, + traffic=ramped(CASE.traffic, ramp, seed=1, concurrency=50), + limits=(Limit("sustained_rps", min=20),), + ) + offsets: Final = send_offsets(ramp, seed=1) + records: Final = tuple( + _record(seq, 1000.0 + offset, 1.5 if ramp.step_of(seq) == 0 else 4.0) + for seq, offset in enumerate(offsets, start=1) + ) + observation: Final = Observation( + records=records, + mock_requests=(), + guardrail_calls=(), + rss_before=NO_RSS, + rss_after=NO_RSS, + wall_seconds=200.0, + ramp=ramp, + ) + first, second = observation.steps() + assert first.held and first.gateway_added_p99 == 0.5 and first.completed_rps > 9.5 + assert not second.held and second.gateway_added_p99 == 3.0 + rows, exceeded = evaluate(case, observation) + assert exceeded == ["sustained_rps: observed 10.0000 < floor 20.0000"] + assert {("perf_sustained_rps", 10.0), ("perf_sustained_rps_floor", 20.0)} <= set(junit_properties(rows)) + assert dict(step_properties((first, second))) == { + "perf_ramp_hold_p99_seconds": 2.0, + "perf_ramp_step_10_gateway_added_p99": 0.5, + "perf_ramp_step_10_completed_rps": first.completed_rps, + "perf_ramp_step_10_unexpected": 0.0, + "perf_ramp_step_10_gateway_added_p50": 0.5, + "perf_ramp_step_20_gateway_added_p99": 3.0, + "perf_ramp_step_20_completed_rps": second.completed_rps, + "perf_ramp_step_20_unexpected": 0.0, + "perf_ramp_step_20_gateway_added_p50": 3.0, + } + assert not evaluate(replace(case, limits=(Limit("sustained_rps", min=10),)), observation)[1] + + +def test_each_step_reads_the_pod_s_cpu_redis_latency_and_connections_from_its_own_window() -> None: + ramp: Final = Ramp(first_rps=10.0, last_rps=20.0, step_rps=10.0, step_seconds=100.0, held_p99_seconds=2.0) + records: Final = tuple( + _record(seq, 1000.0 + offset, 1.5) for seq, offset in enumerate(send_offsets(ramp, seed=1), start=1) + ) + started: Final = min(record.started_at for record in records) + redis: Final = tuple( + Sample(at=started + at, value=GatewayTelemetry(redis_latency=HistogramTotals(total=total, count=count))) + for at, total, count in ((-1.0, 0.0, 0.0), (50.0, 1.0, 100.0), (99.0, 2.0, 200.0), (199.0, 42.0, 300.0)) + ) + observation: Final = Observation( + records=records, + mock_requests=(), + guardrail_calls=(), + rss_before=NO_RSS, + rss_after=NO_RSS, + wall_seconds=200.0, + ramp=ramp, + cpu_series=Series(samples=tuple(Sample(at=started + at, value=5.0) for at in (10.0, 60.0, 110.0, 160.0))), + connection_series=Series( + samples=tuple( + Sample(at=started + at, value=n) for at, n in ((1.0, 12), (50.0, 18), (150.0, 25), (199.0, 20)) + ) + ), + telemetry_series=Series(samples=redis), + ) + first, second = observation.steps() + first_completed: Final = round(first.completed_rps * ramp.step_seconds) + assert first.cpu_seconds_per_request is not None and math.isclose( + first.cpu_seconds_per_request, 10.0 / first_completed + ) + assert first.redis_latency_mean_seconds is not None and math.isclose(first.redis_latency_mean_seconds, 0.01) + assert first.postgres_peak_connections == 18 + assert second.redis_latency_mean_seconds is not None and math.isclose(second.redis_latency_mean_seconds, 0.4) + assert second.postgres_peak_connections == 25 + assert first.time_to_first_token_p99 is None + names: Final = {name for name, _ in step_properties((first, second)) if name.startswith("perf_ramp_step_10_")} + assert names == { + "perf_ramp_step_10_gateway_added_p99", + "perf_ramp_step_10_completed_rps", + "perf_ramp_step_10_unexpected", + "perf_ramp_step_10_gateway_added_p50", + "perf_ramp_step_10_cpu_seconds_per_request", + "perf_ramp_step_10_redis_latency_mean_seconds", + "perf_ramp_step_10_postgres_peak_connections", + } diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py index 9a6c9d2af14..3414476ba4c 100644 --- a/tests/e2e_harness/perf/test_upgrade.py +++ b/tests/e2e_harness/perf/test_upgrade.py @@ -304,7 +304,9 @@ def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None: def _ceiling(metric: str, percentile: int | None = None) -> float: return next( - limit.max for limit in upgrade_case().limits if limit.metric == metric and limit.percentile == percentile + limit.max + for limit in upgrade_case().limits + if limit.metric == metric and limit.percentile == percentile and limit.max is not None ) diff --git a/tests/integration/authorization/test_claude_code_gateway_oauth.py b/tests/integration/authorization/test_claude_code_gateway_oauth.py index 384bf3a2918..a80eeacad46 100644 --- a/tests/integration/authorization/test_claude_code_gateway_oauth.py +++ b/tests/integration/authorization/test_claude_code_gateway_oauth.py @@ -15,6 +15,7 @@ from urllib.parse import parse_qs, urlencode, urlsplit import httpx import pytest import yaml +from anthropic import Anthropic from anthropic.types import Message from integration._support import claude_code as cc from integration._support.client import Gateway, eventually, gateway_from_environment @@ -397,7 +398,31 @@ def test_device_authorization_polls_complete_sso_and_mints_a_model_restricted_be "anthropic-beta": cc.CLI_BETA, "x-claude-code-session-id": f"oauth-{user_id}", } - request_body: Final = {**request_template, "model": model_a} + standard_listing: Final = enabled_gateway.client.get("/v1/models", headers=bearer_headers) + gateway_listing: Final = enabled_gateway.client.get( + "/claude_code_gateway/v1/models", headers=bearer_headers + ) + assert standard_listing.status_code == 200, standard_listing.text + assert gateway_listing.status_code == 200, gateway_listing.text + assert gateway_listing.json() == standard_listing.json() + with Anthropic( + base_url=f"{str(enabled_gateway.client.base_url).rstrip('/')}/claude_code_gateway", + auth_token=access_token.access_token, + ) as sdk: + discovered_models: Final = tuple(sdk.models.list()) + assert {model.id for model in discovered_models} == {model_a, model_c} + discovered_model: Final = next(model for model in discovered_models if model.id == model_a) + openai_listing: Final = enabled_gateway.client.get( + "/claude_code_gateway/v1/models", + headers={"authorization": f"Bearer {access_token.access_token}"}, + ) + assert openai_listing.status_code == 200, openai_listing.text + assert openai_listing.json()["object"] == "list" + assert {model["id"] for model in openai_listing.json()["data"]} == {model_a, model_c} + unauthenticated: Final = enabled_gateway.client.get("/claude_code_gateway/v1/models") + assert unauthenticated.status_code == 401, unauthenticated.text + assert upstream.drain() == () + request_body: Final = {**request_template, "model": discovered_model.id} expected_message: Final = _MESSAGE.validate_python( { **_JSON_OBJECT.validate_json( diff --git a/tests/integration/compatibility/test_openai_consumer.py b/tests/integration/compatibility/test_openai_consumer.py index a49f20f4711..6150995723b 100644 --- a/tests/integration/compatibility/test_openai_consumer.py +++ b/tests/integration/compatibility/test_openai_consumer.py @@ -459,7 +459,6 @@ async def test_openai_clients_parse_text_completion_streams_and_choices(gateway: def test_text_completion_usage_chunk_has_empty_choices(gateway: Gateway) -> None: - pytest.skip("BUG: the /v1/completions include_usage chunk carries choices [{index: 0}] instead of []") identity: Final = f"cmpl-usage-{uuid.uuid4().hex}" prompt: Final = f"raw-stream {identity}" with wire_server(_text_provider(identity, {prompt})) as wire, gateway.scenario() as scenario: diff --git a/tests/integration/compatibility/test_status_code_followups_usage_chat.py b/tests/integration/compatibility/test_status_code_followups_usage_chat.py deleted file mode 100644 index de9ec16cf30..00000000000 --- a/tests/integration/compatibility/test_status_code_followups_usage_chat.py +++ /dev/null @@ -1,125 +0,0 @@ -"""The /usage/ai/chat tools hand the model the reason the daily activity query rejected its dates, the same check -the daily activity routes answer 400 with, instead of a generic fetch error.""" - -from __future__ import annotations - -import json -import uuid -from collections.abc import Iterator -from dataclasses import dataclass -from typing import Final - -import pytest -from integration._support.client import JSON_OBJECT, Gateway, object_value -from integration._support.upstream import delete_scenario, register_scenario -from integration.compatibility._status_code_audit import Upstream, json_response, owned_gateway -from pydantic import JsonValue - -pytestmark: Final = pytest.mark.timeout(180) - -_TOOLS: Final = ("get_usage_data", "get_team_usage_data", "get_tag_usage_data") - - -@dataclass(frozen=True, slots=True) -class _UsageChat: - gateway: Gateway - identity: str - - -@pytest.fixture(scope="module") -def usage_chat(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_UsageChat]: - """A proxy whose usage chat model (`openai/...`, called straight through litellm) is the scripted upstream.""" - identity: Final = f"audit-usage-chat-{uuid.uuid4().hex}" - handle: Final = register_scenario(identity, json_response({})) - try: - with owned_gateway( - tmp_path_factory.mktemp("audit-usage-chat"), - {"model_list": [], "router_settings": {"num_retries": 0}}, - environment={"OPENAI_API_BASE": handle.api_base(), "OPENAI_API_KEY": identity}, - ) as owned: - yield _UsageChat(owned, identity) - finally: - delete_scenario(handle) - - -def _tool_calls(start_date: str, end_date: str) -> dict[str, JsonValue]: - arguments: Final = json.dumps({"start_date": start_date, "end_date": end_date}) - return { - "id": "chatcmpl-usage", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": None, - "tool_calls": [ - {"id": f"call-{tool}", "type": "function", "function": {"name": tool, "arguments": arguments}} - for tool in _TOOLS - ], - }, - "finish_reason": "tool_calls", - } - ], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - } - - -def _tool_results(usage_chat: _UsageChat, start_date: str, end_date: str) -> tuple[dict[str, str], tuple[str, ...]]: - """Runs one chat turn whose model calls every tool with these dates; returns the tool messages the follow-up - model call carried, keyed by tool, and the tool statuses the SSE stream reported.""" - register_scenario(usage_chat.identity, json_response(_tool_calls(start_date, end_date))) - Upstream(usage_chat.gateway.upstream_url).reset() - with usage_chat.gateway.client.stream( - "POST", - "/usage/ai/chat", - json={"model": "openai/gpt-4o-mini", "messages": [{"role": "user", "content": "How much did we spend?"}]}, - headers={"Authorization": f"Bearer {usage_chat.gateway.key}"}, - timeout=60, - ) as response: - assert response.status_code == 200 - events: Final = tuple( - object_value(JSON_OBJECT.validate_json(line.removeprefix("data: "))) - for line in response.iter_lines() - if line.startswith("data: ") - ) - calls: Final = Upstream(usage_chat.gateway.upstream_url).calls(usage_chat.identity) - assert len(calls) == 2, calls - messages: Final = object_value(calls[1]["body"]).get("messages") - assert isinstance(messages, list), calls[1] - tool_messages: Final = { - str(message["tool_call_id"]).removeprefix("call-"): str(message["content"]) - for message in messages - if isinstance(message, dict) and message.get("role") == "tool" - } - statuses: Final = tuple(str(event.get("status")) for event in events if event.get("type") == "tool_call") - return tool_messages, statuses - - -_REVERSED: Final = "Invalid date range: end_date must be on or after start_date" -_NOT_CANONICAL: Final = "Invalid date range: start_date and end_date must be valid YYYY-MM-DD dates" -_NO_TEAM_DATA: Final = "No Team usage data found for the given date range." -_NO_TAG_DATA: Final = "No Tag usage data found for the given date range." - - -@pytest.mark.parametrize( - ("start_date", "end_date", "team_message", "tag_message"), - ( - pytest.param("2026-10-09", "2026-10-01", _REVERSED, _REVERSED, id="reversed"), - pytest.param("2026-13-01", "2026-10-09", _NOT_CANONICAL, _NOT_CANONICAL, id="invalid-month"), - pytest.param("2001/01/01", "2001-01-09", _NOT_CANONICAL, _NOT_CANONICAL, id="slashes"), - pytest.param("2001-01-01", "2001-01-09", _NO_TEAM_DATA, _NO_TAG_DATA, id="valid"), - ), -) -def test_usage_chat_tools_answer_with_the_query_layers_date_check( - usage_chat: _UsageChat, start_date: str, end_date: str, team_message: str, tag_message: str -) -> None: - """The team and tag tools give the model the daily activity routes' 400 reason; the global tool, whose query - takes these dates without that check, still answers with its usage summary.""" - tool_messages, statuses = _tool_results(usage_chat, start_date, end_date) - assert tool_messages["get_usage_data"].startswith("Total Spend: $"), tool_messages - assert tool_messages["get_team_usage_data"] == team_message, tool_messages - assert tool_messages["get_tag_usage_data"] == tag_message, tool_messages - assert statuses == ("running", "complete") * len(_TOOLS), statuses diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index acba66223e5..53fd5c1cec2 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -6985,6 +6985,106 @@ GPT_5_6_RESPONSES_CACHE_READ: Final = CostTrackingTestCase( ), ) +GPT_5_6_RESPONSES_CACHE_WRITE_12H: Final = CostTrackingTestCase( + name="gpt-5.6-responses_cache_write_12h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + endpoint="/v1/responses", + request={ + "model": "$MODEL", + "input": "summarize this text", + "prompt_cache_options": {"mode": "implicit", "ttl": "12h"}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + } + ], + "usage": { + "input_tokens": 12928, + "output_tokens": 380, + "total_tokens": 13308, + "input_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 5000, + "cache_write_tokens_details": {"ttl_30m_tokens": 3000, "ttl_12h_tokens": 2000}, + }, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + ), + expected=ExactExpected( + spend=0.0344865, + input_cost=0.0291865, + output_cost=0.00532, + prompt_tokens=12928, + completion_tokens=380, + cache_creation_cost=0.0153125, + ), +) + +GPT_6_LUNA_RESPONSES_CACHE_WRITE_12H_NO_12H_RATE: Final = CostTrackingTestCase( + name="gpt-6-luna-responses_cache_write_12h_no_12h_rate", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-luna", + endpoint="/v1/responses", + request={ + "model": "$MODEL", + "input": "summarize this text", + "prompt_cache_options": {"mode": "implicit", "ttl": "12h"}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6-luna", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + } + ], + "usage": { + "input_tokens": 12928, + "output_tokens": 380, + "total_tokens": 13308, + "input_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 5000, + "cache_write_tokens_details": {"ttl_30m_tokens": 3000, "ttl_12h_tokens": 2000}, + }, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + ), + expected=ExactExpected( + spend=0.0016078, + input_cost=0.0014178, + output_cost=0.00019, + prompt_tokens=12928, + completion_tokens=380, + cache_creation_cost=0.000625, + ), +) + GPT_5_6_RESPONSES_REASONING: Final = CostTrackingTestCase( name="gpt-5.6-responses_reasoning", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -11913,6 +12013,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( OMNI_MODERATIONS_NEXT_SINGLE, OMNI_MODERATIONS_NEXT_LIST, GPT_5_6_RESPONSES_CACHE_READ, + GPT_5_6_RESPONSES_CACHE_WRITE_12H, + GPT_6_LUNA_RESPONSES_CACHE_WRITE_12H_NO_12H_RATE, GPT_5_6_RESPONSES_REASONING, GPT_5_6_RESPONSES_STREAM, GPT_5_6_RESPONSES_STREAM_CACHE_READ, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 012d847e676..cd979528fe6 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -174,6 +174,8 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { }, }, "gpt-5.6": { + "cache_creation_input_token_cost": 2.1875e-06, + "cache_creation_input_token_cost_above_12hr": 4.375e-06, "cache_read_input_token_cost": 1.75e-07, "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 1.75e-06, diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index e98f6a44222..95363e92366 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -85,6 +85,16 @@ class CostMapEntry(BaseModel): cache_read_input_token_cost: float | None = None cache_creation_input_token_cost: float | None = None cache_creation_input_token_cost_above_1hr: float | None = None + cache_creation_input_token_cost_above_12hr: float | None = None + cache_creation_input_token_cost_above_12hr_flex: float | None = None + cache_creation_input_token_cost_above_12hr_priority: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority: float | None = None + cache_creation_input_token_cost_above_12hr_batches: float | None = None + cache_creation_input_token_cost_above_12hr_ultrafast: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast: float | None = None cache_creation_input_token_cost_above_1hr_above_200k_tokens: float | None = None cache_read_input_token_cost_above_200k_tokens: float | None = None cache_creation_input_token_cost_above_200k_tokens: float | None = None diff --git a/tests/integration/database/test_cleanup_transactions.py b/tests/integration/database/test_cleanup_transactions.py index 72da2bb48ee..202657d180e 100644 --- a/tests/integration/database/test_cleanup_transactions.py +++ b/tests/integration/database/test_cleanup_transactions.py @@ -113,7 +113,7 @@ async def test_size_rule_deletes_the_oldest_rows_until_under_the_target_and_does """ Deleting rows never shrinks the table file, so the size pass must stop on live bytes, not on pg_total_relation_size, taking the oldest rows first and keeping the tool index on the same cutoff. Once - autovacuum has turned the deleted rows into free space the catalog no longer shows them as dead, and the + VACUUM has turned the deleted rows into free space the catalog no longer shows them as dead, and the files still measure the same; the next check must read what the last run freed from its saved count and leave the table alone instead of calling the whole file live again. """ @@ -140,13 +140,15 @@ async def test_size_rule_deletes_the_oldest_rows_until_under_the_target_and_does setup.execute( sql.SQL( 'CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL, payload text, ' - 'PRIMARY KEY (request_id, "startTime"))' + 'PRIMARY KEY (request_id, "startTime")) ' + 'WITH (autovacuum_enabled = false, toast.autovacuum_enabled = false)' ).format(logs) ) setup.execute( sql.SQL( "CREATE TABLE {} (request_id text, tool_name text, start_time timestamp NOT NULL, " - "PRIMARY KEY (request_id, tool_name))" + "PRIMARY KEY (request_id, tool_name)) " + "WITH (autovacuum_enabled = false, toast.autovacuum_enabled = false)" ).format(index) ) setup.execute(sql.SQL("CREATE TABLE {} (param_name text PRIMARY KEY, param_value jsonb)").format(config)) diff --git a/tests/integration/management/test_budget_list_sorting.py b/tests/integration/management/test_budget_list_sorting.py new file mode 100644 index 00000000000..aae0bd6e274 --- /dev/null +++ b/tests/integration/management/test_budget_list_sorting.py @@ -0,0 +1,94 @@ +from datetime import datetime, timezone +from itertools import chain +from typing import Final +from uuid import uuid4 + +import pytest + +from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from +from tests.integration._support.client import Gateway, Scenario + + +@pytest.mark.parametrize("descending", (False, True)) +def test_budget_reset_sort_orders_durations_before_pagination( + descending: bool, + gateway: Gateway, + scenario: Scenario, +) -> None: + prefix: Final = f"sort-{uuid4().hex}" + durations: Final = ("30d", "7d", "1h", "60m", "1mo", "weekly", "monthly", "٥d", "۵d", "५d", "5d", "𝟝d", "5d", None) + ids: Final = tuple( + scenario.budget(budget_id=f"{prefix}-{index}", budget_duration=duration, max_budget=10) + for index, duration in enumerate(durations) + ) + now: Final = datetime.now(timezone.utc) + expected: Final = tuple( + budget_id + for _, budget_id in sorted( + zip(durations, ids, strict=True), + key=lambda item: ( + item[0] is None, + duration_in_seconds_from(now, item[0]) * (-1 if descending else 1) if item[0] else 0, + item[1], + ), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/management/v1/budgets", + params={ + "q": prefix, + "sort": "-budget_duration" if descending else "budget_duration", + "page_size": "3", + "page": str(page), + }, + ) + for page in range(1, -(-len(ids) // 3) + 1) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert tuple(row["budget_id"] for row in chain.from_iterable(page.json()["data"] for page in pages)) == expected + filtered: Final = gateway.request( + "GET", + "/management/v1/budgets", + params={ + "q": prefix, + "sort": "budget_duration", + "filter[budget_duration][in]": "7d,30d", + }, + ) + assert filtered.status_code == 200, filtered.text + assert [row["budget_duration"] for row in filtered.json()["data"]] == ["7d", "30d"] + + +@pytest.mark.parametrize("field", ("max_budget", "tpm_limit", "rpm_limit", "tpd_limit")) +@pytest.mark.parametrize("descending", (False, True)) +def test_budget_limits_sort_unlimited_as_infinity( + field: str, descending: bool, gateway: Gateway, scenario: Scenario +) -> None: + prefix: Final = f"sort-unlimited-{uuid4().hex}" + values: Final = (20, None, 0, 3) + ids: Final = tuple(scenario.budget(budget_id=f"{prefix}-{i}", **{field: value}) for i, value in enumerate(values)) + expected: Final = tuple( + identity + for value, identity in sorted( + zip(values, ids, strict=True), + key=lambda item: float("inf") if item[0] is None else item[0], + reverse=descending, + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/management/v1/budgets", + params={ + "q": prefix, + "sort": ("-" if descending else "") + field, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert tuple(row["budget_id"] for row in chain.from_iterable(page.json()["data"] for page in pages)) == expected diff --git a/tests/integration/management/test_key_list.py b/tests/integration/management/test_key_list.py index b70a3782eda..4634314817b 100644 --- a/tests/integration/management/test_key_list.py +++ b/tests/integration/management/test_key_list.py @@ -1,11 +1,16 @@ -from typing import FrozenSet +from datetime import datetime, timedelta, timezone +from itertools import chain +from typing import Final, FrozenSet import pytest from litellm.proxy.utils import hash_token -from tests.integration._support.actors import TEAM_ALPHA, Actor +from tests.integration._support.actors import TEAM_ALPHA, Actor, ManagementDatabase, ScratchResources from tests.integration._support.actors import create_scratch_key +from tests.integration._support.client import Gateway, Scenario +from litellm.repositories.table_repositories import DeletedVerificationTokenRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -202,3 +207,74 @@ async def test_key_list_non_admin_cannot_filter_other_team(gateway_client, actor }, ) assert resp.status_code == 403, resp.text + + +@pytest.mark.parametrize( + "field", ("expires", "last_active", "budget_reset_at", "total_spend", "tpm_limit", "rpm_limit") +) +@pytest.mark.parametrize("direction", ("asc", "desc")) +@pytest.mark.parametrize("deleted", (False, True), ids=("live", "deleted")) +async def test_key_list_sorts_across_pages_with_nulls_last( + field: str, + direction: str, + deleted: bool, + gateway: Gateway, + scenario: Scenario, + db: ManagementDatabase, + resources: ScratchResources, +) -> None: + team_admin: Final = scenario.user(user_role="internal_user") + owner: Final = scenario.user(user_role="internal_user") + team_id: Final = scenario.team( + members_with_roles=[{"user_id": team_admin, "role": "admin"}, {"user_id": owner, "role": "user"}] + ) + team_admin_key: Final = scenario.key(user_id=team_admin) + table: Final = ( + DeletedVerificationTokenRepository(db.prisma).table if deleted else VerificationTokenRepository(db.prisma).table + ) + values: Final = (20, 3, 3, 100, 0 if field == "total_spend" else None) + tokens: Final = tuple(hash_token(f"{resources.prefix}-{index}") for index in range(len(values))) + rows: Final = [ + { + "token": token, + "key_alias": resources.prefix, + "team_id": team_id, + "user_id": owner, + field: ( + datetime(2030, 1, 1, tzinfo=timezone.utc) + timedelta(days=value) + if field in ("expires", "last_active", "budget_reset_at") and value is not None + else value + ), + } + for token, value in zip(tokens, values, strict=True) + ] + await table.create_many(rows) + try: + valued: Final = tuple((value, token) for value, token in zip(values, tokens, strict=True) if value is not None) + expected: Final = [ + token + for _, token in sorted(valued, key=lambda item: (item[0] if direction == "asc" else -item[0], item[1])) + ] + [token for value, token in zip(values, tokens, strict=True) if value is None] + for caller in (gateway.key, team_admin_key): + pages: Final = [ + gateway.client.get( + "/key/list", + headers={"Authorization": f"Bearer {caller}"}, + params={ + "key_alias": resources.prefix, + "team_id": team_id, + "include_team_keys": "true", + "sort_by": field, + "sort_order": direction, + "size": "2", + "page": str(page), + **({"status": "deleted"} if deleted else {}), + }, + ) + for page in (1, 2, 3) + ] + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert all(page.json()["total_count"] == len(tokens) for page in pages) + assert list(chain.from_iterable(page.json()["keys"] for page in pages)) == expected + finally: + await table.delete_many(where={"token": {"in": tokens}}) diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index c0c03571a13..d0b6e2579c8 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/tests/integration/management/test_sso_sign_in.py @@ -3,7 +3,9 @@ import contextlib import json import uuid from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from functools import partial from pathlib import Path from typing import Final from urllib.parse import parse_qs, urlencode, urlparse @@ -15,7 +17,7 @@ from pydantic import JsonValue from tests.integration._support.client import Gateway, gateway_from_environment from tests.integration._support.database import read_rows -from tests.integration._support.process import owned_proxy +from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState, owned_proxy from tests.integration._support.wire import Reply, Request, wire_server pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] @@ -128,3 +130,38 @@ def test_a_first_sso_sign_in_creates_the_user_with_the_default_role( "metadata": {"auth_provider": IDP_PROVIDER}, } ] + + +def test_free_sso_admission_preserves_existing_users_and_serializes_the_last_seat( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv(OWNED_STATE_VARIABLE, OwnedState.PRIVATE.value) + monkeypatch.delenv("LITELLM_LICENSE", raising=False) + with _sso_gateway(tmp_path, {}) as gateway: + for index in range(4): + signed_in: Final = _sign_in(gateway, f"free-sso-{index}") + assert signed_in.status_code == 303, signed_in.text + + contenders: Final = tuple(f"free-sso-contender-{index}" for index in range(8)) + with ThreadPoolExecutor(max_workers=len(contenders)) as pool: + responses: Final = tuple(pool.map(partial(_sign_in, gateway), contenders)) + + assert sorted(response.status_code for response in responses) == [303, 403, 403, 403, 403, 403, 403, 403] + assert gateway.rows('SELECT COUNT(*) AS count FROM "LiteLLM_UserTable"', ()) == [{"count": 5}] + for subject, response in zip(contenders, responses): + if response.status_code == 403: + assert "requires a LiteLLM Enterprise license" in response.text + assert not response.cookies + assert gateway.rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) == [] + + existing: Final = _sign_in(gateway, "free-sso-0") + assert existing.status_code == 303, existing.text + + admin_created: Final = gateway.request( + "POST", "/user/new", {"user_id": "api-created-user", "auto_create_key": False} + ) + assert admin_created.status_code == 200, admin_created.text + above_capacity: Final = _sign_in(gateway, "free-sso-0") + refused: Final = _sign_in(gateway, "free-sso-extra") + assert (above_capacity.status_code, refused.status_code) == (303, 403) + assert gateway.rows('SELECT COUNT(*) AS count FROM "LiteLLM_UserTable"', ()) == [{"count": 6}] diff --git a/tests/integration/management/test_team_list_sorting.py b/tests/integration/management/test_team_list_sorting.py new file mode 100644 index 00000000000..c09e51f0dfe --- /dev/null +++ b/tests/integration/management/test_team_list_sorting.py @@ -0,0 +1,244 @@ +import json +from itertools import chain +from typing import Final +from uuid import uuid4 + +import pytest + +from tests.integration._support.client import Gateway, Scenario +from tests.integration._support.database import write_rows + + +@pytest.mark.parametrize( + "field", ("spend", "max_budget", "tpm_limit", "rpm_limit", "updated_at", "members", "models", "keys_count") +) +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_team_list_sorts_across_pages_and_respects_membership( + field: str, + direction: str, + gateway: Gateway, + scenario: Scenario, +) -> None: + prefix: Final = f"sort-{uuid4().hex}" + owner: Final = scenario.user(user_role="internal_user") + caller: Final = scenario.key(user_id=owner) + values: Final = (20, 3, 3, 100, 0) + teams: Final = tuple( + scenario.team( + team_id=f"{prefix}-{index}", + team_alias=f"{prefix}-{index}", + members_with_roles=[{"user_id": owner, "role": "admin"}], + ) + for index in range(len(values)) + ) + outsider: Final = scenario.team(team_alias=f"{prefix}-outside") + for team_id, value in zip(teams, values, strict=True): + write_rows( + 'UPDATE "LiteLLM_TeamTable" SET spend = %s::float, max_budget = %s::float, ' + "tpm_limit = %s::bigint, rpm_limit = %s::bigint, " + "updated_at = timestamp '2030-01-01' + %s::int * interval '1 day', " + "members_with_roles = %s::jsonb, models = ARRAY(SELECT generate_series(1, %s::int)::text) " + "WHERE team_id = %s", + ( + str(value), + str(value), + str(value), + str(value), + str(value), + json.dumps( + [{"user_id": owner if i == 0 else f"{prefix}-member-{i}", "role": "admin"} for i in range(value)] + ), + str(value), + team_id, + ), + ) + write_rows( + 'INSERT INTO "LiteLLM_VerificationToken" (token, team_id) ' + "SELECT %s || '-key-' || n::text, %s FROM generate_series(1, %s::int) n", + (team_id, team_id, str(value)), + ) + write_rows( + 'UPDATE "LiteLLM_TeamTable" SET max_budget = NULL, tpm_limit = NULL, rpm_limit = NULL WHERE team_id = %s', + (teams[-1],), + ) + try: + expected: Final = tuple( + team_id + for value, team_id in sorted( + zip(values, teams, strict=True), + key=lambda item: ( + (item[0] != 0 if direction == "desc" else item[0] == 0) + if field in ("max_budget", "tpm_limit", "rpm_limit") + else False, + item[0] * (1 if direction == "asc" else -1), + item[1], + ), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/v2/team/list", + key=caller, + params={ + "search": prefix, + "user_id": owner, + "sort_by": field, + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2, 3) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + returned: Final = tuple(row["team_id"] for row in chain.from_iterable(page.json()["teams"] for page in pages)) + assert returned == expected + assert outsider not in returned + assert all(page.json()["total"] == len(teams) for page in pages) + finally: + for team_id in teams: + write_rows('DELETE FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team_id,)) + + +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_team_list_sorts_organization_names_and_missing_values( + direction: str, gateway: Gateway, scenario: Scenario +) -> None: + prefix: Final = f"sort-org-{uuid4().hex}" + budget: Final = scenario.budget(max_budget=10) + organizations: Final = ((f"{prefix}-z", "alpha"), (f"{prefix}-a", "zulu"), (f"{prefix}-m", "")) + teams: Final = tuple(scenario.team(team_id=f"{prefix}-{i}", team_alias=f"{prefix}-{i}") for i in range(4)) + try: + for (organization_id, alias), team_id in zip(organizations, teams): + write_rows( + 'INSERT INTO "LiteLLM_OrganizationTable" ' + "(organization_id, organization_alias, budget_id, created_by, updated_by, updated_at) " + "VALUES (%s, %s, %s, 'sorting-test', 'sorting-test', NOW())", + (organization_id, alias, budget), + ) + write_rows( + 'UPDATE "LiteLLM_TeamTable" SET organization_id = %s WHERE team_id = %s', + (organization_id, team_id), + ) + ascending: Final = (teams[0], teams[2], teams[1]) + ordered: Final = ascending if direction == "asc" else tuple(reversed(ascending)) + pages: Final = tuple( + gateway.request( + "GET", + "/v2/team/list", + params={ + "search": prefix, + "sort_by": "organization_alias", + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert tuple(row["team_id"] for row in chain.from_iterable(page.json()["teams"] for page in pages)) == (*ordered, teams[-1]) + finally: + for team_id in teams: + write_rows('UPDATE "LiteLLM_TeamTable" SET organization_id = NULL WHERE team_id = %s', (team_id,)) + for organization_id, _ in organizations: + write_rows('DELETE FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (organization_id,)) + + +@pytest.mark.parametrize("field", ("organization_alias", "members", "models", "keys_count")) +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_computed_team_sorts_preserve_org_and_membership_union( + field: str, direction: str, gateway: Gateway, scenario: Scenario +) -> None: + prefix: Final = f"sort-org-scope-{uuid4().hex}" + owner: Final = scenario.user(user_role="internal_user") + budget: Final = scenario.budget(max_budget=10) + orgs: Final = (f"{prefix}-allowed", f"{prefix}-outside") + teams: Final = tuple( + scenario.team( + team_id=f"{prefix}-{i}", + team_alias=f"{prefix}-{i}", + members_with_roles=[{"user_id": owner, "role": "admin"}] if i == 2 else [], + ) + for i in range(4) + ) + try: + for org, alias in zip(orgs, ("alpha", "zulu"), strict=True): + write_rows( + 'INSERT INTO "LiteLLM_OrganizationTable" ' + "(organization_id, organization_alias, budget_id, created_by, updated_by, updated_at) " + "VALUES (%s, %s, %s, 'sorting-test', 'sorting-test', NOW())", + (org, alias, budget), + ) + write_rows( + 'INSERT INTO "LiteLLM_OrganizationMembership" (user_id, organization_id, user_role) ' + "VALUES (%s, %s, 'org_admin')", + (owner, orgs[0]), + ) + for index, team in enumerate(teams): + write_rows( + 'UPDATE "LiteLLM_TeamTable" SET organization_id = %s, members_with_roles = %s::jsonb, ' + "models = ARRAY(SELECT generate_series(1, %s::int)::text) WHERE team_id = %s", + ( + orgs[0] if index < 2 else orgs[1], + json.dumps( + [ + {"user_id": owner if i == 0 else f"{prefix}-member-{i}", "role": "admin"} + for i in range(index + 1) + ] + ), + str(index + 1), + team, + ), + ) + write_rows( + 'INSERT INTO "LiteLLM_VerificationToken" (token, team_id) ' + "SELECT %s || '-key-' || n::text, %s FROM generate_series(1, %s::int) n", + (team, team, str(index + 1)), + ) + caller: Final = scenario.key(user_id=owner) + values: Final = (0, 0, 1) if field == "organization_alias" else (1, 2, 3) + expected: Final = tuple( + team + for value, team in sorted( + zip(values, teams[:3], strict=True), + key=lambda item: (item[0] * (-1 if direction == "desc" else 1), item[1]), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/v2/team/list", + key=caller, + params={ + "user_id": owner, + "search": prefix, + "sort_by": field, + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert all(page.json()["total"] == 3 for page in pages) + assert tuple(row["team_id"] for row in chain.from_iterable(page.json()["teams"] for page in pages)) == expected + denied: Final = gateway.request( + "GET", + "/v2/team/list", + key=caller, + params={ + "organization_id": orgs[1], + "sort_by": field, + }, + ) + assert denied.status_code == 403, denied.text + finally: + for team in teams: + write_rows('DELETE FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,)) + write_rows('UPDATE "LiteLLM_TeamTable" SET organization_id = NULL WHERE team_id = %s', (team,)) + for org in orgs: + write_rows('DELETE FROM "LiteLLM_OrganizationMembership" WHERE organization_id = %s', (org,)) + write_rows('DELETE FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (org,)) diff --git a/tests/integration/management/test_team_list_v2.py b/tests/integration/management/test_team_list_v2.py index d608d054daa..4c9171310a6 100644 --- a/tests/integration/management/test_team_list_v2.py +++ b/tests/integration/management/test_team_list_v2.py @@ -1,8 +1,8 @@ -from typing import FrozenSet, Optional +from typing import Final, FrozenSet, Optional import pytest -from tests.integration._support.actors import Actor +from tests.integration._support.actors import Actor, ActorSet, GatewayClient pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -139,3 +139,31 @@ async def test_team_list_v2_invalid_status_is_400(gateway_client, actors): headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"}, ) assert resp.status_code == 400, resp.text + + +async def test_team_list_v2_returns_caller_membership_only_for_members( + gateway_client: GatewayClient, + actors: ActorSet, +) -> None: + member: Final = actors.keys[Actor.TEAM_ADMIN] + member_response: Final = await gateway_client.get( + f"/v2/team/list?team_id={actors.team_alpha_id}&user_id={member.user_id}", + headers={"Authorization": f"Bearer {member.cleartext}"}, + ) + + assert member_response.status_code == 200, member_response.text + member_team: Final = member_response.json()["teams"][0] + assert member_team["caller_membership"] == { + "spend": 0.0, + "max_budget": None, + "budget_reset_at": None, + } + + admin: Final = actors.keys[Actor.PROXY_ADMIN] + admin_response: Final = await gateway_client.get( + f"/v2/team/list?team_id={actors.team_alpha_id}", + headers={"Authorization": f"Bearer {admin.cleartext}"}, + ) + + assert admin_response.status_code == 200, admin_response.text + assert admin_response.json()["teams"][0]["caller_membership"] is None diff --git a/tests/integration/management/test_user_list_sorting.py b/tests/integration/management/test_user_list_sorting.py new file mode 100644 index 00000000000..f5578140e1c --- /dev/null +++ b/tests/integration/management/test_user_list_sorting.py @@ -0,0 +1,179 @@ +from itertools import chain +from typing import Final +from uuid import uuid4 + +import pytest + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from tests.integration._support.client import Gateway, Scenario +from tests.integration._support.database import write_rows + + +@pytest.mark.parametrize("field", ("user_alias", "max_budget", "updated_at", "sso_user_id", "status", "key_count")) +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_user_list_sorts_all_pages(field: str, direction: str, gateway: Gateway, scenario: Scenario) -> None: + prefix: Final = f"sort-{uuid4().hex}" + ids: Final = tuple( + scenario.user( + user_id=f"{prefix}-{index}", + user_role="internal_user", + user_alias=f"Alias {value if index != 2 else 4:03}", + user_email=f"{prefix}-{index}@example.com", + max_budget=value, + metadata={"scim_active": value != 3}, + ) + for index, value in enumerate((20, 3, 3, 100, 0)) + ) + for user_id, value in zip(ids, (20, 3, 3, 100, 0), strict=True): + write_rows( + 'UPDATE "LiteLLM_UserTable" SET sso_user_id = %s, ' + "updated_at = timestamp '2030-01-01' + %s::int * interval '1 day' WHERE user_id = %s", + (f"{prefix}-{value:03}", str(value), user_id), + ) if value != 3 or user_id == ids[1] else write_rows( + "UPDATE \"LiteLLM_UserTable\" SET sso_user_id = %s, updated_at = timestamp '2030-01-04' WHERE user_id = %s", + (f"{prefix}-004", user_id), + ) + for index in range(value if value < 4 else 1): + scenario.key(user_id=user_id, key_alias=f"{user_id}-key-{index}") + write_rows( + 'UPDATE "LiteLLM_UserTable" SET max_budget = NULL, updated_at = NULL, sso_user_id = NULL WHERE user_id = %s', + (ids[-1],), + ) + scenario.key(user_id=ids[-1]) + write_rows( + 'UPDATE "LiteLLM_VerificationToken" SET team_id = %s WHERE user_id = %s', (UI_SESSION_TOKEN_TEAM_ID, ids[-1]) + ) + values: Final = ( + (0, 1, 1, 0, 0) + if field == "status" + else (1, 3, 3, 1, 0) + if field == "key_count" + else (20, 3, 4, 100, None) + if field == "sso_user_id" + else (20, 3, 4, 100, 0) + if field == "user_alias" + else (20, 3, 3, 100, None) + ) + expected: Final = tuple( + user_id + for _, user_id in sorted( + ((value, user_id) for value, user_id in zip(values, ids, strict=True)), + key=lambda item: ( + (item[0] is not None if field == "max_budget" and direction == "desc" else item[0] is None), + (item[0] or 0) * (1 if direction == "asc" else -1), + item[1], + ), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/user/list", + params={ + "search": prefix, + "sort_by": field, + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2, 3) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert all(page.json()["total"] == len(ids) for page in pages) + assert tuple(user["user_id"] for user in chain.from_iterable(page.json()["users"] for page in pages)) == expected + + +@pytest.mark.parametrize("field", ("status", "key_count")) +def test_computed_user_sorts_preserve_filters(field: str, gateway: Gateway, scenario: Scenario) -> None: + team: Final = scenario.team() + member: Final = scenario.member(team) + other: Final = scenario.user(user_role="internal_user") + for filters in ({"team": team}, {"user_ids": member}, {"user_ids": f"{member},{other}"}, {"role": "internal_user"}): + baseline: Final = gateway.get("/user/list", {"user_ids": f"{member},{other}", **filters, "sort_by": "user_id"}) + response: Final = gateway.get("/user/list", {"user_ids": f"{member},{other}", **filters, "sort_by": field}) + assert response["total"] == baseline["total"] + assert sorted(user["user_id"] for user in response["users"]) == sorted( + user["user_id"] for user in baseline["users"] + ) + ordinary_key: Final = scenario.key(user_id=other) + denied: Final = gateway.request("GET", "/user/list", key=ordinary_key, params={"sort_by": field}) + assert denied.status_code == 403, denied.text + + +@pytest.mark.parametrize("field", ("status", "key_count")) +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_computed_user_sorts_preserve_org_admin_scope( + field: str, direction: str, gateway: Gateway, scenario: Scenario +) -> None: + prefix: Final = f"sort-org-users-{uuid4().hex}" + owner: Final = scenario.user(user_role="internal_user") + members: Final = tuple( + scenario.user(user_id=f"{prefix}-{index}", user_role="internal_user", metadata={"scim_active": index != 1}) + for index in range(4) + ) + budget: Final = scenario.budget(max_budget=10) + orgs: Final = (f"{prefix}-allowed", f"{prefix}-outside") + try: + for org in orgs: + write_rows( + 'INSERT INTO "LiteLLM_OrganizationTable" ' + "(organization_id, organization_alias, budget_id, created_by, updated_by, updated_at) " + "VALUES (%s, %s, %s, 'sorting-test', 'sorting-test', NOW())", + (org, org, budget), + ) + write_rows( + 'INSERT INTO "LiteLLM_OrganizationMembership" (user_id, organization_id, user_role) ' + "VALUES (%s, %s, 'org_admin')", + (owner, orgs[0]), + ) + for index, member in enumerate(members): + write_rows( + 'INSERT INTO "LiteLLM_OrganizationMembership" (user_id, organization_id, user_role) ' + "VALUES (%s, %s, 'internal_user')", + (member, orgs[0] if index < 3 else orgs[1]), + ) + for _ in range(index): + scenario.key(user_id=member) + caller: Final = scenario.key(user_id=owner) + values: Final = (0, 1, 0) if field == "status" else (0, 1, 2) + expected: Final = tuple( + member + for value, member in sorted( + zip(values, members[:3], strict=True), + key=lambda item: (item[0] * (-1 if direction == "desc" else 1), item[1]), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/user/list", + key=caller, + params={ + "search": prefix, + "sort_by": field, + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert all(page.json()["total"] == 3 for page in pages) + returned: Final = tuple(row["user_id"] for row in chain.from_iterable(page.json()["users"] for page in pages)) + assert returned == expected + denied: Final = gateway.request( + "GET", + "/user/list", + key=caller, + params={ + "organization_ids": orgs[1], + "sort_by": field, + }, + ) + assert denied.status_code == 403, denied.text + finally: + for org in orgs: + write_rows('DELETE FROM "LiteLLM_OrganizationMembership" WHERE organization_id = %s', (org,)) + write_rows('DELETE FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (org,)) diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 327b053d333..90a462472b2 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -792,17 +792,23 @@ def _ui_session_cookie(gateway: Gateway, user_id: str) -> dict[str, str]: def _gateway_session_sign_in( - gateway: Gateway, user_id: str, resource: str | None = None + gateway: Gateway, + user_id: str, + resource: str | None = None, + *, + registration_endpoint: str = "/register", + authorization_endpoint: str = "/authorize", + token_endpoint: str = "/token", ) -> tuple[dict[str, object], str]: registered: Final = gateway.client.post( - "/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"} + registration_endpoint, json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"} ) assert registered.status_code in (200, 201), registered.text client_id: Final = str(_response_object(registered)["client_id"]) pkce: Final = _Pkce(secrets.token_urlsafe(48)) cookies: Final = _ui_session_cookie(gateway, user_id) started: Final = gateway.client.get( - "/authorize", + authorization_endpoint, params={ "client_id": client_id, "redirect_uri": CLIENT_REDIRECT, @@ -823,13 +829,14 @@ def _gateway_session_sign_in( callback: Final = parse_qs(urlsplit(completed.headers["location"]).query) assert "code" in callback, completed.headers["location"] issued: Final = gateway.client.post( - "/token", + token_endpoint, data={ "grant_type": "authorization_code", "code": callback["code"][0], "redirect_uri": CLIENT_REDIRECT, "client_id": client_id, "code_verifier": pkce.verifier, + **({} if resource is None else {"resource": resource}), }, ) assert issued.status_code == 200, issued.text @@ -1292,10 +1299,18 @@ def test_lite_logout_revokes_the_cli_refresh_token_sent_the_way_the_cli_sends_it assert completion.status_code == 200, completion.text -def _toolset_rpc(gateway: Gateway, bearer: str, name: str, method: str, params: dict[str, object]) -> Outcome: +def _toolset_rpc( + gateway: Gateway, + bearer: str, + name: str, + method: str, + params: dict[str, object], + *, + route: str = "/toolset/{name}/mcp", +) -> Outcome: def post(rpc_method: str, rpc_params: dict[str, object]) -> httpx.Response: return gateway.client.post( - f"/toolset/{name}/mcp", + route.format(name=name), json={"jsonrpc": "2.0", "id": 1, "method": rpc_method, "params": rpc_params}, headers={"Authorization": f"Bearer {bearer}", "Accept": "application/json, text/event-stream"}, ) @@ -1306,7 +1321,12 @@ def _toolset_rpc(gateway: Gateway, bearer: str, name: str, method: str, params: return _outcome_from_rpc(post(method, params)) -def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_its_route(gateway: Gateway) -> None: +@pytest.mark.parametrize("route", ["/toolset/{name}/mcp", "/{name}/mcp"]) +def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_its_route( + gateway: Gateway, route: str +) -> None: + from mcp.shared.auth_utils import check_resource_allowed + with mcp_peer() as peer, gateway.scenario() as scenario: alias: Final = "lit6029sess" + uuid.uuid4().hex[:6] server_id: Final = register_mcp(scenario, peer, alias) @@ -1315,19 +1335,69 @@ def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_i granted_id: Final = create_toolset(scenario, ((server_id, "add"),), toolset_name=granted_name) create_toolset(scenario, ((server_id, "multiply"),), toolset_name=withheld_name) member: Final = scenario.member(scenario.team(object_permission={"mcp_toolsets": [granted_id]})) - bearer: Final = _gateway_session_bearer(gateway, member) + path: Final = route.format(name=granted_name) + anonymous: Final = gateway.client.post( + path, + headers=ACCEPT, + json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": dict(INITIALIZE)}, + ) + assert anonymous.status_code == 401, anonymous.text + metadata_url: Final = anonymous.headers["www-authenticate"].split('resource_metadata="', 1)[1].split('"', 1)[0] + protected: Final = gateway.client.get(metadata_url) + assert protected.status_code == 200, protected.text + protected_body: Final = _response_object(protected) + resource: Final = str(protected_body["resource"]) + assert check_resource_allowed( + requested_resource=f"{_base(gateway)}{path}", configured_resource=resource + ), protected.text + issuers: Final = TypeAdapter(tuple[str, ...]).validate_python(protected_body["authorization_servers"]) + assert len(issuers) == 1, protected.text + issuer: Final = urlsplit(issuers[0]) + authorization: Final = gateway.client.get( + f"{issuer.scheme}://{issuer.netloc}/.well-known/oauth-authorization-server{issuer.path}" + ) + assert authorization.status_code == 200, authorization.text + authorization_body: Final = _response_object(authorization) + assert authorization_body["issuer"] == issuers[0], authorization.text + token_endpoint: Final = str(authorization_body["token_endpoint"]) + issued, client_id = _gateway_session_sign_in( + gateway, + member, + resource, + registration_endpoint=str(authorization_body["registration_endpoint"]), + authorization_endpoint=str(authorization_body["authorization_endpoint"]), + token_endpoint=token_endpoint, + ) + bearer: Final = _issued_token(issued) assert bearer.startswith("llm_session_"), bearer[:16] - listed: Final = _toolset_rpc(gateway, bearer, granted_name, "tools/list", {}) + listed: Final = _toolset_rpc(gateway, bearer, granted_name, "tools/list", {}, route=route) assert listed.tools == (f"{alias}-add",), listed.raw peer.drain() called: Final = _toolset_rpc( - gateway, bearer, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": {"a": 4, "b": 5}} + gateway, bearer, granted_name, "tools/call", + {"name": f"{alias}-add", "arguments": {"a": 4, "b": 5}}, route=route ) assert called.ok and called.text == "9", called.raw assert len(tool_calls(peer.drain())) == 1 - denied: Final = _toolset_rpc(gateway, bearer, withheld_name, "tools/list", {}) + denied: Final = _toolset_rpc(gateway, bearer, withheld_name, "tools/list", {}, route=route) assert denied.status == 403, denied.raw assert tool_calls(peer.drain()) == () + rotated: Final = gateway.client.post( + token_endpoint, + data={ + "grant_type": "refresh_token", + "refresh_token": str(issued["refresh_token"]), + "client_id": client_id, + "resource": resource, + }, + ) + assert rotated.status_code == 200, rotated.text + rotated_body: Final = _response_object(rotated) + assert rotated_body["refresh_token"] != issued["refresh_token"], rotated.text + refreshed: Final = _toolset_rpc( + gateway, _issued_token(rotated_body), granted_name, "tools/list", {}, route=route + ) + assert refreshed.tools == listed.tools, refreshed.raw def test_resource_scoped_session_bearer_opens_a_team_toolset_inside_its_server_and_none_outside( diff --git a/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_disconnect_wire.py b/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_disconnect_wire.py new file mode 100644 index 00000000000..100e760d5be --- /dev/null +++ b/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_disconnect_wire.py @@ -0,0 +1,232 @@ +import asyncio +import json +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import ExitStack, asynccontextmanager +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import httpx +import pytest +import pytest_asyncio +import yaml +from aiohttp import web + +from tests.integration._support.anthropic_sse import LIFECYCLE, delta_text, event_types, parse_sse, user_prompt +from tests.integration._support.client import JSON_OBJECT, Gateway +from tests.integration._support.openai_wire import chat_stream +from tests.integration._support.ports import allocate_port +from tests.integration._support.process import OwnedProxy, owned_proxy_process, paused + +_MODEL: Final = "disconnect-repro" +_KEY: Final = "sk-integration-master-key" +_DEADLINE: Final = 10 +_RACE_ATTEMPTS: Final = 20 + + +class _Mode(Enum): + OCCUPY = "occupy" + QUEUED = "queued" + CONTENT = "content" + COMPLETE = "complete" + + +@dataclass(frozen=True, slots=True) +class _Request: + mode: _Mode + headers_allowed: asyncio.Event = field(default_factory=asyncio.Event) + headers_sent: asyncio.Event = field(default_factory=asyncio.Event) + running: asyncio.Event = field(default_factory=asyncio.Event) + released: asyncio.Event = field(default_factory=asyncio.Event) + + +@dataclass(frozen=True, slots=True) +class _Backend: + url: str + requests: asyncio.Queue[_Request] + runner: web.AppRunner + + +@pytest_asyncio.fixture(loop_scope="function") +async def backend() -> AsyncIterator[_Backend]: + requests: Final[asyncio.Queue[_Request]] = asyncio.Queue() + worker: Final = asyncio.Semaphore(1) + + async def complete(request: web.Request) -> web.StreamResponse: + body: Final = JSON_OBJECT.validate_json(await request.read()) + assert body["model"] == "test-model" and body["stream"] is True, body + assert request.headers["Authorization"] == "Bearer integration-provider-key" + pending: Final = _Request(_Mode(user_prompt(body))) + await requests.put(pending) + response: Final = web.StreamResponse(headers={"Content-Type": "text/event-stream"}) + try: + await pending.headers_allowed.wait() + await response.prepare(request) + pending.headers_sent.set() + async with worker: + pending.running.set() + chunks: Final = chat_stream("chatcmpl-disconnect", "test-model", "hello") + if pending.mode in (_Mode.CONTENT, _Mode.COMPLETE): + for chunk in chunks[:2]: + await response.write(chunk) + if pending.mode is _Mode.COMPLETE: + await response.write(chunks[2]) + await response.write_eof() + else: + await asyncio.Event().wait() + finally: + pending.released.set() + return response + + app: Final = web.Application() + app.router.add_post("/v1/chat/completions", complete) + runner: Final = web.AppRunner(app, handler_cancellation=True, shutdown_timeout=0.1) + await runner.setup() + port: Final = allocate_port() + await web.TCPSite(runner, "127.0.0.1", port).start() + try: + yield _Backend(f"http://127.0.0.1:{port}", requests, runner) + finally: + await runner.cleanup() + + +@pytest_asyncio.fixture(loop_scope="function") +async def proxy(backend: _Backend, tmp_path: Path) -> AsyncIterator[OwnedProxy]: + config: Final = tmp_path / "disconnect.yaml" + config.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _MODEL, + "litellm_params": { + "model": "custom_openai/test-model", + "api_base": backend.url + "/v1", + "api_key": "integration-provider-key", + }, + } + ], + "litellm_settings": {"anthropic_sse_ping_interval_seconds": 0}, + "router_settings": {"num_retries": 0}, + "general_settings": {"dangerously_permit_weak_or_unset_master_key": True}, + } + ) + ) + with httpx.Client() as client: + template: Final = Gateway(client, _KEY, backend.url) + with owned_proxy_process( + template, + tmp_path, + {"STORE_MODEL_IN_DB": "False"}, + config=config, + database_setup=(), + remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "REDIS_HOST", "REDIS_PORT"), + ) as owned: + try: + yield owned + finally: + await backend.runner.cleanup() + + +def _body(mode: _Mode) -> bytes: + return json.dumps( + {"model": _MODEL, "stream": True, "max_tokens": 16, "messages": [{"role": "user", "content": mode.value}]} + ).encode() + + +@asynccontextmanager +async def _request( + proxy: OwnedProxy, mode: _Mode, endpoint: str = "/v1/messages" +) -> AsyncGenerator[asyncio.StreamReader, None]: + address: Final = urlsplit(str(proxy.gateway.client.base_url)) + reader, writer = await asyncio.open_connection(address.hostname, address.port) + body: Final = _body(mode) + writer.write( + f"POST {endpoint} HTTP/1.1\r\nHost: localhost\r\nAuthorization: Bearer {_KEY}\r\n" + f"Content-Type: application/json\r\nContent-Length: {len(body)}\r\n\r\n".encode() + + body + ) + await writer.drain() + try: + yield reader + finally: + writer.close() + await writer.wait_closed() + + +async def _next(backend: _Backend, mode: _Mode) -> _Request: + pending: Final = await asyncio.wait_for(backend.requests.get(), _DEADLINE) + assert pending.mode is mode + return pending + + +async def _headers(pending: _Request) -> None: + pending.headers_allowed.set() + await asyncio.wait_for(pending.headers_sent.wait(), _DEADLINE) + + +async def _released(pending: _Request) -> None: + try: + await asyncio.wait_for(pending.released.wait(), _DEADLINE) + except TimeoutError: + pytest.fail(f"Backend {pending.mode.value} request still active after client disconnect", pytrace=False) + + +async def _disconnect_queued_request(backend: _Backend, proxy: OwnedProxy) -> None: + with ExitStack() as scheduling: + async with _request(proxy, _Mode.QUEUED): + queued: Final = await _next(backend, _Mode.QUEUED) + scheduling.enter_context(paused(proxy.process)) + await _headers(queued) + assert not queued.running.is_set() + await _released(queued) + assert backend.requests.empty(), "Disconnect must not trigger another upstream request" + + +async def test_disconnect_with_upstream_headers_pending_releases_queued_request( + backend: _Backend, proxy: OwnedProxy +) -> None: + """Hold headers and the client FIN ready together without patching proxy code.""" + async with _request(proxy, _Mode.OCCUPY): + occupied: Final = await _next(backend, _Mode.OCCUPY) + await _headers(occupied) + await asyncio.wait_for(occupied.running.wait(), _DEADLINE) + for attempt in range(_RACE_ATTEMPTS): + await _disconnect_queued_request(backend, proxy) + assert not occupied.released.is_set(), attempt + await _released(occupied) + + +@pytest.mark.parametrize("endpoint", ["/v1/messages", "/v1/chat/completions"]) +async def test_disconnect_after_content_releases_upstream(backend: _Backend, proxy: OwnedProxy, endpoint: str) -> None: + async with _request(proxy, _Mode.CONTENT, endpoint) as response: + pending: Final = await _next(backend, _Mode.CONTENT) + await _headers(pending) + received: Final = await asyncio.wait_for(response.readuntil(b'"hello"'), _DEADLINE) + assert b"200 OK" in received + if endpoint == "/v1/messages": + assert b"content_block_delta" in received + await _released(pending) + assert backend.requests.empty(), "Disconnect must not trigger another upstream request" + + +async def test_messages_stream_completes_with_unchanged_events(backend: _Backend, proxy: OwnedProxy) -> None: + async with httpx.AsyncClient(base_url=str(proxy.gateway.client.base_url), trust_env=False) as client: + response_task: Final = asyncio.create_task( + client.post( + "/v1/messages", + headers={"Authorization": f"Bearer {_KEY}", "Content-Type": "application/json"}, + content=_body(_Mode.COMPLETE), + ) + ) + pending: Final = await _next(backend, _Mode.COMPLETE) + await _headers(pending) + response: Final = await asyncio.wait_for(response_task, _DEADLINE) + assert response.status_code == 200, response.text + events: Final = parse_sse(response.text) + assert event_types(events) == LIFECYCLE + assert delta_text(events) == "hello" + assert pending.released.is_set() + assert backend.requests.empty(), "Successful streaming must not trigger another upstream request" diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index e0b2e4a336f..041a0a72fdd 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -1,13 +1,17 @@ import json import os import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager from dataclasses import dataclass +from pathlib import Path from typing import Final, Literal import anthropic import httpx +from pydantic import JsonValue, TypeAdapter import pytest +import yaml from integration._support.anthropic_sse import ( LIFECYCLE, Attempts, @@ -27,13 +31,14 @@ from integration._support.anthropic_sse import ( ) from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server -from pydantic import JsonValue from redis import Redis _MODEL: Final = "claude-sonnet-4-5-20250929" _API_KEY: Final = "synthetic-anthropic-key" _TEXT: Final = "Hello" +_SINK_EVENT_BATCH_ADAPTER: Final = TypeAdapter(list[dict[str, JsonValue]]) FirstAttempt = Literal[ "drop_after_headers", @@ -127,24 +132,24 @@ def _stream(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[int, tuple return response.status_code, parse_sse(response.text) -def _rows(request_id: str) -> list[dict[str, JsonValue]]: - return read_rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)) +def _rows(gateway: Gateway, request_id: str) -> list[dict[str, JsonValue]]: + return gateway.rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)) -def _success_rows(request_id: str) -> list[dict[str, JsonValue]]: - return eventually(lambda: _rows(request_id), lambda rows: len(rows) >= 1, seconds=70) +def _success_rows(gateway: Gateway, request_id: str) -> list[dict[str, JsonValue]]: + return eventually(lambda: _rows(gateway, request_id), lambda rows: len(rows) >= 1, seconds=70) def _assert_completed_by_retry( - status: int, events: tuple[SseEvent, ...], prompt: str, model: str, wire: Wire, *, attempts: int = 2 + gateway: Gateway, status: int, events: tuple[SseEvent, ...], prompt: str, model: str, wire: Wire, *, attempts: int = 2 ) -> None: assert status == 200, events assert event_types(events) == LIFECYCLE, events assert message_id(events) == _served_id(prompt, attempts), events assert delta_text(events) == _TEXT, events assert [request.target for request in wire.drain()] == ["/v1/messages"] * attempts - assert _success_rows(_served_id(prompt, attempts)) == [{"status": "success", "model_group": model}] - superseded_rows: Final = {attempt: _rows(_served_id(prompt, attempt)) for attempt in range(1, attempts)} + assert _success_rows(gateway, _served_id(prompt, attempts)) == [{"status": "success", "model_group": model}] + superseded_rows: Final = {attempt: _rows(gateway, _served_id(prompt, attempt)) for attempt in range(1, attempts)} assert not any(superseded_rows.values()), superseded_rows @@ -183,7 +188,7 @@ def test_stream_failing_before_content_is_retried_on_the_same_group_and_complete **({"num_retries": 1} if source == "deployment" else {}), ) status, events = _stream(gateway, _body(model, prompt, source)) - _assert_completed_by_retry(status, events, prompt, model, wire) + _assert_completed_by_retry(gateway, status, events, prompt, model, wire) @pytest.mark.parametrize("source", ["deployment", "request"]) @@ -204,7 +209,7 @@ def test_stream_rejected_before_it_opens_is_retried_and_completes( **({"num_retries": 1} if source == "deployment" else {}), ) status, events = _stream(gateway, _body(model, prompt, source)) - _assert_completed_by_retry(status, events, prompt, model, wire) + _assert_completed_by_retry(gateway, status, events, prompt, model, wire) def _sdk(gateway: Gateway) -> anthropic.Anthropic: @@ -232,7 +237,7 @@ def test_anthropic_sdk_sync_stream_completes_after_a_drop_following_message_star assert final.id == _served_id(prompt, 2), final assert [block.text for block in final.content if block.type == "text"] == [_TEXT], final assert [request.target for request in wire.drain()] == ["/v1/messages"] * 2 - assert _success_rows(final.id) == [{"status": "success", "model_group": model}] + assert _success_rows(gateway, final.id) == [{"status": "success", "model_group": model}] async def test_anthropic_sdk_async_stream_completes_after_a_drop_following_message_start(gateway: Gateway) -> None: @@ -321,7 +326,7 @@ def test_stream_error_frame_after_content_is_not_retried_and_keeps_its_spend_row assert message_id(events) == _served_id(prompt, 1), events assert delta_text(events) == _TEXT, events assert [request.target for request in wire.drain()] == ["/v1/messages"] - assert [row["model_group"] for row in _success_rows(_served_id(prompt, 1))] == [model] + assert [row["model_group"] for row in _success_rows(gateway, _served_id(prompt, 1))] == [model] def test_stream_rejected_with_401_before_it_opens_is_not_retried(gateway: Gateway) -> None: @@ -393,6 +398,22 @@ def test_retry_rejected_before_it_opens_counts_against_the_same_budget(gateway: model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) response: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "deployment")) _assert_rejected_before_content(response, wire, status=500, attempts=2, error="error") + rows: Final = eventually( + lambda: _request_rows(gateway, response.headers["x-litellm-call-id"]), + lambda values: len(values) >= 1, + seconds=70, + ) + rejected: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "request", budget=0)) + _assert_rejected_before_content(rejected, wire, status=500, attempts=1, error="error") + rejection_rows: Final = eventually( + lambda: _request_rows(gateway, rejected.headers["x-litellm-call-id"]), + lambda values: len(values) >= 1, + seconds=70, + ) + estimated_input: Final = rejection_rows[0]["prompt_tokens"] + assert isinstance(estimated_input, int) + assert rejection_rows == [{"status": "failure", "prompt_tokens": estimated_input, "completion_tokens": 0}] + assert rows == [{"status": "failure", "prompt_tokens": 5 + estimated_input, "completion_tokens": 1}], rows def test_non_stream_messages_rejected_with_500_is_retried_per_the_deployment_budget(gateway: Gateway) -> None: @@ -409,7 +430,7 @@ def test_non_stream_messages_rejected_with_500_is_retried_per_the_deployment_bud assert payload["id"] == _served_id(prompt, 2), response.text assert payload["content"] == [{"type": "text", "text": _TEXT}], response.text assert [request.target for request in wire.drain()] == ["/v1/messages"] * 2 - assert _success_rows(_served_id(prompt, 2)) == [{"status": "success", "model_group": model}] + assert _success_rows(gateway, _served_id(prompt, 2)) == [{"status": "success", "model_group": model}] def test_retried_stream_response_headers_name_the_attempt(gateway: Gateway) -> None: @@ -422,7 +443,7 @@ def test_retried_stream_response_headers_name_the_attempt(gateway: Gateway) -> N model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) response: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "deployment")) events: Final = parse_sse(response.text) - _assert_completed_by_retry(response.status_code, events, prompt, model, wire) + _assert_completed_by_retry(gateway, response.status_code, events, prompt, model, wire) assert response.headers.get("x-litellm-attempted-retries") == "1", dict(response.headers) assert response.headers.get("x-litellm-max-retries") == "1", dict(response.headers) @@ -454,7 +475,7 @@ def test_retried_stream_spend_row_records_the_attempt_count(gateway: Gateway) -> ): model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) status, events = _stream(gateway, _body(model, prompt, "deployment")) - _assert_completed_by_retry(status, events, prompt, model, wire) + _assert_completed_by_retry(gateway, status, events, prompt, model, wire) rows: Final = read_rows( "SELECT metadata->>'attempted_retries' AS attempted, metadata->>'max_retries' AS budget " 'FROM "LiteLLM_SpendLogs" WHERE request_id=%s', @@ -463,6 +484,114 @@ def test_retried_stream_spend_row_records_the_attempt_count(gateway: Gateway) -> assert rows == [{"attempted": "1", "budget": "1"}], rows +def _proxy_with_callback_sink(gateway: Gateway, tmp_path: Path, sink: Wire) -> AbstractContextManager[Gateway]: + loaded_config: Final = TypeAdapter(dict[str, object]).validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + litellm_settings: Final = TypeAdapter(dict[str, object]).validate_python(loaded_config["litellm_settings"]) + config: Final = { + **loaded_config, + "litellm_settings": { + **litellm_settings, + "callbacks": ["generic_api"], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + }, + } + path: Final = tmp_path / "callbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": sink.url, "GENERIC_LOGGER_HEADERS": "Authorization=Bearer synthetic-sink-secret"}, + config=path, + ) + + +def _sink_events(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]: + for batch in batches: + yield from _SINK_EVENT_BATCH_ADAPTER.validate_python(json.loads(batch.body)) + + +class _SinkReader: + """Reads the callback sink across polls: drain() consumes the sink's queue, so earlier batches are kept here.""" + + def __init__(self, sink: Wire, model: str) -> None: + self._sink: Final = sink + self._model: Final = model + self._batches: tuple[Request, ...] = () + + def delivered_usage(self) -> tuple[tuple[str, int, int], ...]: + self._batches = (*self._batches, *self._sink.drain()) + return tuple( + (str(event["status"]), int(event["prompt_tokens"]), int(event["completion_tokens"])) + for event in _sink_events(self._batches) + if event.get("model_group") == self._model + ) + + +def _request_rows(gateway: Gateway, call_id: str) -> list[dict[str, JsonValue]]: + return gateway.rows( + 'SELECT status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,) + ) + + +@pytest.mark.parametrize( + "failure", + (_Failure("drop_after_message_start"), _Failure("error_frame_after_message_start", 500)), + ids=("transport-drop", "recoverable-error-frame"), +) +def test_superseded_attempt_and_its_retry_each_log_their_own_usage( + gateway: Gateway, tmp_path: Path, failure: _Failure +) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, failure, attempts)) as wire, + wire_server(lambda request: Reply()) as sink, + _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + status, events = _stream(candidate, _body(model, prompt, "deployment")) + _assert_completed_by_retry(candidate, status, events, prompt, model, wire) + delivered: Final = eventually( + _SinkReader(sink, model).delivered_usage, lambda values: len(values) >= 2, seconds=20 + ) + assert sorted(delivered) == [("failure", 5, 1), ("success", 5, 3)], delivered + + +def test_every_dropped_attempt_logs_its_own_failure_and_the_request_row_still_bills_the_dropped_usage( + gateway: Gateway, tmp_path: Path +) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, _Failure("drop_after_message_start"), attempts, failing_attempts=99)) as wire, + wire_server(lambda request: Reply()) as sink, + _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + response: Final = candidate.request("POST", "/v1/messages", _body(model, prompt, "deployment")) + _assert_stream_failed_after_message_start(response.status_code, parse_sse(response.text), wire, attempts=2) + delivered: Final = eventually( + _SinkReader(sink, model).delivered_usage, lambda values: len(values) >= 2, seconds=20 + ) + assert delivered == (("failure", 5, 1), ("failure", 5, 1)), delivered + rows: Final = eventually( + lambda: _request_rows(candidate, response.headers["x-litellm-call-id"]), + lambda values: len(values) >= 1, + seconds=70, + ) + assert rows == [ + { + "status": "failure", + "prompt_tokens": sum(input_tokens for _, input_tokens, _ in delivered), + "completion_tokens": sum(output_tokens for _, _, output_tokens in delivered), + } + ], rows + + def _string_values(cache: Redis) -> tuple[bytes, ...]: keys: Final = tuple(key for key in cache.scan_iter(count=1000) if cache.type(key) == b"string") return tuple(value for value in cache.mget(keys) if value is not None) if keys else () @@ -485,7 +614,7 @@ def test_retried_stream_is_cached_and_the_identical_request_is_served_without_th model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) body: Final = _body(model, prompt, "deployment") status, events = _stream(gateway, body) - _assert_completed_by_retry(status, events, prompt, model, wire) + _assert_completed_by_retry(gateway, status, events, prompt, model, wire) eventually(lambda: _cached_somewhere(_served_id(prompt, 2)), bool, seconds=30) replay_status, replay = _stream(gateway, body) assert replay_status == 200, replay diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index b46503703c7..58785c89033 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -312,7 +312,7 @@ def _chat_caller_stream_usage_event(event: dict[str, JsonValue], model: str) -> "object": event["object"], "created": 1, "model": model, - "choices": [{"index": 0, "delta": {}}], + "choices": [], "usage": {**_json_object_value(event["usage"]), "cost": 4.05e-6}, } diff --git a/tests/integration/observability/test_otel_team_sampling_across_pods.py b/tests/integration/observability/test_otel_team_sampling_across_pods.py new file mode 100644 index 00000000000..dc7b295febd --- /dev/null +++ b/tests/integration/observability/test_otel_team_sampling_across_pods.py @@ -0,0 +1,723 @@ +"""A team's Arize sampling rate decides once per trace, so every pod and worker serving a call of it agrees. + +Two proxies share one database, as two pods behind a load balancer do: the first runs one worker, the second +two. Calls of one trace carry the caller's ``traceparent`` and land on different pods and workers; the team's +Arize space must receive every call of a trace or none of them. +""" + +import asyncio +import json +import re +import uuid +from collections.abc import Callable, Generator, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass, replace +from itertools import chain, product +from pathlib import Path +from typing import Final, Literal, TypeAlias +from urllib.parse import urlparse + +import httpx +import pytest +import yaml +from integration._support.client import ( + Scenario, + eventually, + gateway_from_environment, + list_value, + object_value, + until_asserted, +) +from integration._support.database import scratch_database +from integration._support.otlp_sink import Span, SpanSinks, configure_sink, recorded_requests, recorded_spans +from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +MODEL: Final = "team-sampling" +SERVER: Final = 2 +CALL_ID: Final = "litellm.call_id" +OPERATION: Final = "gen_ai.operation.name" +SPACE: Final = "space-sampled" +CREDENTIALS: Final = {"arize_space_key": SPACE, "arize_api_key": "arize-sampled"} +FAILING: Final = re.compile(rb"sampling-fail") +JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +TRACES: Final = 16 +Endpoint: TypeAlias = Literal["chat", "responses", "messages"] +Verdicts: TypeAlias = Mapping[str, tuple[bool, ...]] + + +def _chat_reply(stream: bool) -> Reply: + usage: Final[JsonValue] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9} + if not stream: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-sampling", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": usage, + } + ).encode() + ) + chunk: Final = {"id": "chatcmpl-sampling", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"} + deltas: Final[tuple[JsonValue, ...]] = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**chunk, "choices": [], "usage": usage}, + ) + return Reply( + content_type="text/event-stream", + chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"), + ) + + +def _responses_reply(stream: bool) -> Reply: + response: Final[Mapping[str, JsonValue]] = { + "id": "resp_sampling", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": "msg_sampling", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}, + } + if not stream: + return Reply(body=json.dumps(response).encode()) + events: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_sampling", + "output_index": 0, + "content_index": 0, + "delta": "ok", + }, + {"type": "response.completed", "sequence_number": 2, "response": dict(response)}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), + ) + + +def _upstream(request: Request) -> Reply: + if FAILING.search(request.body) is not None: + return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}') + stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True + if request.target.endswith("/responses"): + return _responses_reply(stream) + return _chat_reply(stream) + + +@dataclass(frozen=True, slots=True) +class Sent: + trace_id: str + parent_span_id: str + call_id: str + + +@dataclass(frozen=True, slots=True) +class Call: + path: str + body: Mapping[str, JsonValue] + headers: Mapping[str, str] + sent: Sent + expected: int + + +def _call( + key: str, trace_id: str | None, endpoint: Endpoint = "chat", stream: bool = False, content: str = "hello" +) -> Call: + parent: Final = uuid.uuid4().hex[:16] + call_id: Final = str(uuid.uuid4()) + messages: Final[JsonValue] = [{"role": "user", "content": content}] + path, body = { + "chat": ("/v1/chat/completions", {"model": MODEL, "messages": messages, "stream": stream}), + "responses": ("/v1/responses", {"model": MODEL, "input": content, "stream": stream}), + "messages": ("/v1/messages", {"model": MODEL, "max_tokens": 16, "messages": messages, "stream": stream}), + }[endpoint] + headers: Final = { + "Authorization": f"Bearer {key}", + "x-litellm-call-id": call_id, + "connection": "close", + **({} if trace_id is None else {"traceparent": f"00-{trace_id}-{parent}-01"}), + } + expected: Final = 500 if FAILING.search(content.encode()) else 200 + return Call(path, body, headers, Sent(trace_id or "", parent if trace_id else "", call_id), expected) + + +@dataclass(frozen=True, slots=True) +class Pods: + """``first`` serves with one worker, ``second`` with two.""" + + first: OwnedProxy + second: OwnedProxy + scenario: Scenario + sinks: SpanSinks + + def team_key(self, variables: Mapping[str, str]) -> str: + team: Final = self.scenario.team() + attached: Final = self.first.gateway.request( + "POST", f"/team/{team}/callback", {"callback_name": "arize", "callback_vars": dict(variables)} + ) + assert attached.status_code == 200, attached.text + return self.scenario.key(team_id=team) + + def send(self, pod: OwnedProxy, call: Call) -> Sent: + with pod.gateway.client.stream("POST", call.path, json=call.body, headers=call.headers) as response: + text: Final = response.read().decode() + assert response.status_code == call.expected, text + return call.sent + + async def send_async(self, client: httpx.AsyncClient, call: Call) -> Sent: + async with client.stream("POST", call.path, json=call.body, headers=call.headers) as response: + text: Final = (await response.aread()).decode() + assert response.status_code == call.expected, text + return call.sent + + def on_each_worker(self, pod: OwnedProxy, call: Callable[[], Call]) -> tuple[Sent, ...]: + """One call served by each worker of ``pod``, the others suspended while it runs.""" + workers: Final = sorted(pod.worker_pids()) + assert len(workers) == pod.gateway.workers, workers + sent: Final[list[Sent]] = [] # mutable-ok: one call per worker, in worker order + for serving in workers: + with _suspended_all_but(pod, workers, serving): + sent.append(self.send(pod, call())) + return tuple(sent) + + def fence(self, key: str) -> tuple[Sent, ...]: + """Calls of an unsampled team on the sampled team's exporter, one per worker of both pods. + + Each worker sends spans to that exporter in the order they were handed to it, so all of these + arriving means every call any worker decided before them has been sent too. + """ + return (self.send(self.first, _call(key, uuid.uuid4().hex)),) + self.on_each_worker( + self.second, lambda: _call(key, uuid.uuid4().hex) + ) + + +@contextmanager +def _suspended_all_but(pod: OwnedProxy, workers: Sequence[int], serving: int) -> Generator[None]: + if len(workers) == 1: + yield + return + others: Final = [pid for pid in workers if pid != serving] + with pod.suspended(others[0]): + yield + + +def _space_id(request: Mapping[str, JsonValue]) -> JsonValue: + return {name.lower(): value for name, value in object_value(request.get("headers") or {}).items()}.get( + "arize-space-id" + ) + + +def _space_span_ids(url: str) -> frozenset[str]: + """Span ids the Arize sink received in export requests sent with the team's space id.""" + ours: Final = ( + list_value(request.get("span_ids") or []) for request in recorded_requests(url) if _space_id(request) == SPACE + ) + return frozenset(span_id for span_id in chain.from_iterable(ours) if isinstance(span_id, str)) + + +def _team_spans(pods: Pods, since: int) -> tuple[Span, ...]: + owned: Final = _space_span_ids(pods.sinks.arize) + return tuple(span for span in recorded_spans(pods.sinks.arize, since)[1] if span["span_id"] in owned) + + +def _delivered(spans: tuple[Span, ...], sent: Sequence[Sent]) -> frozenset[str]: + """Call ids of ``sent`` whose server span is among ``spans``.""" + by_parent: Final = {(call.trace_id, call.parent_span_id): call.call_id for call in sent} + return frozenset( + by_parent[(span["trace_id"], span["parent_span_id"])] + for span in spans + if span["kind"] == SERVER and (span["trace_id"], span["parent_span_id"]) in by_parent + ) + + +def _await_logged(pods: Pods, since: int, sent: Sequence[Sent]) -> None: + """Wait until every call of ``sent`` has its model span in the operator's unsampled copy.""" + expected: Final = {call.call_id for call in sent} + eventually( + lambda: ( + { + str(span["attributes"][CALL_ID]) + for span in recorded_spans(pods.sinks.operator, since)[1] + if OPERATION in span["attributes"] and CALL_ID in span["attributes"] + } + & expected + ), + lambda logged: logged == expected, + seconds=60, + ) + + +@dataclass(frozen=True, slots=True) +class Cursors: + operator: int + arize: int + + +def _cursors(pods: Pods) -> Cursors: + return Cursors(recorded_spans(pods.sinks.operator)[0], recorded_spans(pods.sinks.arize)[0]) + + +def _await_flushed(pods: Pods, cursors: Cursors, fence_key: str) -> None: + """Wait until every worker of both pods has sent the team's exporter everything it decided so far.""" + fences: Final = pods.fence(fence_key) + + def arrived() -> None: + delivered: Final = _delivered(_team_spans(pods, cursors.arize), fences) + assert {fence.call_id for fence in fences} <= delivered, "the unsampled team's calls have not arrived yet" + + until_asserted(arrived, seconds=30) + + +def _settled( + pods: Pods, + cursors: Cursors, + by_trace: Mapping[str, Sequence[Sent]], + fence_key: str, + holds: Callable[[tuple[bool, ...]], bool] = lambda kept: len(set(kept)) == 1, +) -> Verdicts: + """Per trace, which of its calls reached the team's space, once every worker has sent what it decided + and ``holds`` is true of every trace.""" + calls: Final = tuple(chain.from_iterable(by_trace.values())) + _await_logged(pods, cursors.operator, calls) + _await_flushed(pods, cursors, fence_key) + + def check() -> Verdicts: + delivered: Final = _delivered(_team_spans(pods, cursors.arize), calls) + verdicts: Final = { + trace_id: tuple(call.call_id in delivered for call in sent) for trace_id, sent in by_trace.items() + } + broken: Final = {trace_id: kept for trace_id, kept in verdicts.items() if not holds(kept)} + assert broken == {}, f"{len(broken)} of {len(verdicts)} traces were split (per call, kept): {broken}" + return verdicts + + return until_asserted(check, seconds=30) + + +def _trace_ids(cell: str, count: int = TRACES) -> tuple[str, ...]: + """Fixed RFC 4122 shaped trace ids, like the uuid4 hex a caller mints for ``traceparent``.""" + return tuple(uuid.uuid5(uuid.NAMESPACE_URL, f"team-sampling/{cell}/{index}").hex for index in range(count)) + + +def _kept(verdicts: Verdicts) -> int: + return sum(1 for kept in verdicts.values() if all(kept)) + + +@pytest.mark.parametrize( + ("variables", "content"), + [({"arize_success_sampling_rate": "0.5"}, "hello"), ({"arize_error_sampling_rate": "0.5"}, "sampling-fail")], + ids=["success-rate", "error-rate"], +) +def test_each_trace_reaches_the_team_space_whole_or_not_at_all( + pods: Pods, variables: Mapping[str, str], content: str +) -> None: + """The ticket's case: three calls per trace load balanced over the pods.""" + sampled: Final = pods.team_key({**CREDENTIALS, **variables}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple( + pods.send(pod, _call(sampled, trace_id, content=content)) for pod in (pods.first, pods.second, pods.first) + ) + for trace_id in _trace_ids("-".join(sorted(variables))) + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES, f"a rate of 0.5 kept {_kept(verdicts)} of {TRACES} traces" + + +@pytest.mark.parametrize("client", ["sync", "async"]) +def test_every_endpoint_follows_the_verdict_of_its_trace(pods: Pods, client: str) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + plan: Final[tuple[tuple[OwnedProxy, Endpoint, bool], ...]] = ( + (pods.first, "chat", False), + (pods.second, "chat", True), + (pods.first, "responses", False), + (pods.second, "messages", True), + ) + + async def send_trace(trace_id: str) -> tuple[Sent, ...]: + clients: Final = { + id(pod): httpx.AsyncClient(base_url=str(pod.gateway.client.base_url), timeout=15, trust_env=False) + for pod in (pods.first, pods.second) + } + try: + return tuple( + await asyncio.gather( + *( + pods.send_async(clients[id(pod)], _call(sampled, trace_id, endpoint, stream)) + for pod, endpoint, stream in plan + ) + ) + ) + finally: + for opened in clients.values(): + await opened.aclose() + + by_trace: Final = { + trace_id: ( + tuple(pods.send(pod, _call(sampled, trace_id, endpoint, stream)) for pod, endpoint, stream in plan) + if client == "sync" + else asyncio.run(send_trace(trace_id)) + ) + for trace_id in _trace_ids(f"endpoints-{client}") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_the_workers_of_one_pod_agree(pods: Pods) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: pods.on_each_worker(pods.second, lambda trace_id=trace_id: _call(sampled, trace_id)) + for trace_id in _trace_ids("workers") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_replaced_worker_reaches_the_verdict_its_predecessor_reached(pods: Pods) -> None: + """A worker that restarts forgets the verdicts it remembered; the calls it serves next for the same + traces still answer the way the first ones did.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + fence_key: Final = pods.team_key(CREDENTIALS) + traces: Final = _trace_ids("replaced") + before: Final = { + trace_id: pods.on_each_worker(pods.second, lambda t=trace_id: _call(sampled, t)) for trace_id in traces + } + _await_logged(pods, cursors.operator, tuple(chain.from_iterable(before.values()))) + _await_flushed(pods, cursors, fence_key) + for victim in sorted(pods.second.worker_pids()): + pods.second.kill_worker(victim) + by_trace: Final = { + trace_id: before[trace_id] + pods.on_each_worker(pods.second, lambda t=trace_id: _call(sampled, t)) + for trace_id in traces + } + + verdicts: Final = _settled(pods, cursors, by_trace, fence_key) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_failed_call_is_kept_wherever_a_successful_call_of_its_trace_is(pods: Pods) -> None: + """One call of the trace fails on one pod and answers to the error rate, another succeeds on the other + pod and answers to the success rate: with the error rate above the success rate, no trace keeps its + successful call and drops its failed one.""" + sampled: Final = pods.team_key( + {**CREDENTIALS, "arize_success_sampling_rate": "0.3", "arize_error_sampling_rate": "0.7"} + ) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: ( + pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")), + pods.send(pods.second, _call(sampled, trace_id)), + ) + for trace_id in _trace_ids("mixed", 32) + } + + verdicts: Final = _settled( + pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept[0] or not kept[1] + ) + + assert any(kept == (True, False) for kept in verdicts.values()), "no trace sat between the two rates" + assert any(kept == (True, True) for kept in verdicts.values()), "no trace sat under both rates" + + +def test_a_failed_call_answers_to_the_error_rate_on_a_pod_that_dropped_its_trace(pods: Pods) -> None: + """A pod that already dropped a successful call of the trace still keeps a later failed call of it at an + unset error rate, as the pod that only served the failed call does.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.3"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("failed-after-dropped") + succeeded: Final = {trace_id: pods.send(pods.first, _call(sampled, trace_id)) for trace_id in traces} + _await_logged(pods, cursors.operator, tuple(succeeded.values())) + by_trace: Final = { + trace_id: ( + succeeded[trace_id], + pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")), + pods.send(pods.second, _call(sampled, trace_id, content="sampling-fail")), + ) + for trace_id in traces + } + + verdicts: Final = _settled( + pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept[1] and kept[2] + ) + + assert any(not kept[0] for kept in verdicts.values()), "no successful call was dropped before its trace failed" + + +def test_a_successful_call_answers_to_the_success_rate_on_a_pod_that_kept_its_failed_call(pods: Pods) -> None: + """A pod that already kept a failed call of the trace drops a later successful call of it at a zero + success rate, as the pod that only served the successful call does.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.0"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("succeeded-after-kept") + failed: Final = { + trace_id: pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")) for trace_id in traces + } + _await_logged(pods, cursors.operator, tuple(failed.values())) + by_trace: Final = { + trace_id: ( + failed[trace_id], + pods.send(pods.first, _call(sampled, trace_id)), + pods.send(pods.second, _call(sampled, trace_id)), + ) + for trace_id in traces + } + + _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept == (True, False, False)) + + +def test_concurrent_calls_of_one_trace_on_one_worker_each_answer_to_their_own_outcome(pods: Pods) -> None: + """A failed and a successful call of the trace in flight at once on the single worker of the first pod + are decided apart: the failed one is kept at an unset error rate, the successful one dropped at a zero + success rate.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.0"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("concurrent-mixed") + plan: Final = tuple(product(traces, ("sampling-fail", "hello"))) + + def send(job: tuple[str, str]) -> Sent: + return pods.send(pods.first, _call(sampled, job[0], content=job[1])) + + with ThreadPoolExecutor(max_workers=8) as pool: + sent: Final = tuple(pool.map(send, plan)) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept == (True, False)) + + +@pytest.mark.parametrize( + ("variables", "expected"), + [({"arize_success_sampling_rate": "1.0"}, TRACES), ({"arize_success_sampling_rate": "0.0"}, 0), ({}, TRACES)], + ids=["one", "zero", "unset"], +) +def test_a_rate_at_either_end_keeps_every_trace_or_none( + pods: Pods, variables: Mapping[str, str], expected: int +) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, **variables}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple(pods.send(pod, _call(sampled, trace_id)) for pod in (pods.first, pods.second)) + for trace_id in _trace_ids(f"edge-{'-'.join(f'{k}={v}' for k, v in sorted(variables.items()))}") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert _kept(verdicts) == expected + + +def test_a_request_without_a_traceparent_is_sampled_whole_at_the_rate(pods: Pods) -> None: + """With no caller trace the proxy starts one per request: each request's server and model spans travel + together and the rate still keeps some requests and drops others.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + sent: Final = tuple( + pods.send(pod, _call(sampled, None)) for _, pod in product(range(12), (pods.first, pods.second)) + ) + _await_logged(pods, cursors.operator, sent) + trace_of: Final = { + str(span["attributes"][CALL_ID]): span["trace_id"] + for span in recorded_spans(pods.sinks.operator, cursors.operator)[1] + if CALL_ID in span["attributes"] + } + fences: Final = pods.fence(pods.team_key(CREDENTIALS)) + + def check() -> Mapping[str, tuple[bool, bool]]: + spans: Final = _team_spans(pods, cursors.arize) + assert {fence.call_id for fence in fences} <= _delivered(spans, fences), ( + "the unsampled team's calls have not arrived yet" + ) + kinds: Final = { + call.call_id: ( + any(span["kind"] == SERVER for span in spans if span["trace_id"] == trace_of[call.call_id]), + any(OPERATION in span["attributes"] for span in spans if span["trace_id"] == trace_of[call.call_id]), + ) + for call in sent + } + partial: Final = {call_id: kept for call_id, kept in kinds.items() if kept[0] != kept[1]} + assert partial == {}, f"requests reached the space without their whole tree (server, model): {partial}" + return kinds + + kept: Final = sum(1 for server, _ in until_asserted(check, seconds=30).values() if server) + assert 0 < kept < len(sent), f"a rate of 0.5 kept {kept} of {len(sent)} requests" + + +def test_a_key_level_destination_follows_the_verdict_of_its_trace(pods: Pods) -> None: + team: Final = pods.scenario.team() + sampled: Final = pods.scenario.key( + team_id=team, + metadata={ + "logging": [ + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {**CREDENTIALS, "arize_success_sampling_rate": "0.5"}, + } + ] + }, + ) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple(pods.send(pod, _call(sampled, trace_id)) for pod in (pods.first, pods.second, pods.first)) + for trace_id in _trace_ids("key-level") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_concurrent_burst_keeps_every_trace_whole(pods: Pods) -> None: + """Four calls per trace, all in flight at once across both pods.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("burst", 24) + plan: Final = tuple(product(traces, (pods.first, pods.second, pods.first, pods.second))) + + def send(job: tuple[str, OwnedProxy]) -> Sent: + return pods.send(job[1], _call(sampled, job[0])) + + with ThreadPoolExecutor(max_workers=16) as pool: + sent: Final = list(pool.map(send, plan)) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < len(traces) + + +def test_a_stalled_team_space_during_a_burst_keeps_serving_and_each_trace_whole(pods: Pods) -> None: + """The team's Arize endpoint stops answering while a burst of calls runs over both pods: every call is + still served, and once the endpoint answers again each trace arrives whole or not at all.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("stalled") + plan: Final = tuple(product(traces, (pods.first, pods.second))) + + def send(job: tuple[str, OwnedProxy]) -> Sent: + return pods.send(job[1], _call(sampled, job[0])) + + configure_sink(pods.sinks.arize, paused=True) + try: + with ThreadPoolExecutor(max_workers=8) as pool: + sent: Final = tuple(pool.map(send, plan)) + finally: + configure_sink(pods.sinks.arize, paused=False) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < len(traces) + + +def _write_config(directory: Path, sinks: SpanSinks, provider: Wire) -> Path: + shared: Final = object_value( + JSON.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + ) + config: Final = { + **shared, + "model_list": [ + *list_value(shared.get("model_list") or []), + { + "model_name": MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": provider.url + "/v1", + }, + }, + ], + "litellm_settings": { + **object_value(shared.get("litellm_settings") or {}), + "callbacks": ["langfuse_otel"], + "otel_tenant_destination_mode": "additive", + "provider_url_destination_allowed_hosts": [urlparse(sinks.operator).netloc, urlparse(sinks.arize).netloc], + }, + "router_settings": {**object_value(shared.get("router_settings") or {}), "num_retries": 0}, + "general_settings": {**object_value(shared.get("general_settings") or {}), "user_api_key_cache_ttl": 2}, + } + path: Final = directory / "team-sampling.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@contextmanager +def _two_pods(provider: Wire, sinks: SpanSinks, directory: Path) -> Generator[Pods]: + environment: Final = { + "LITELLM_OTEL_V2": "1", + "OTEL_BSP_SCHEDULE_DELAY": "300", + "MAPPER_NAMES": "genai,openinference", + "LANGFUSE_HOST": sinks.operator, + "LANGFUSE_PUBLIC_KEY": "pk-lf-operator", + "LANGFUSE_SECRET_KEY": "sk-lf-operator", + "ARIZE_SPACE_KEY": "space-operator", + "ARIZE_API_KEY": "arize-operator", + "ARIZE_HTTP_ENDPOINT": sinks.arize + "/v1/traces", + } + removed: Final = ("LITELLM_OTEL_TENANT_DESTINATION_MODE", "OTEL_TRACES_SAMPLER", "OTEL_TRACES_SAMPLER_ARG") + config: Final = _write_config(directory, sinks, provider) + for pod in ("first", "second"): + (directory / pod).mkdir() + with gateway_from_environment() as shared, scratch_database("team_sampling") as database_url: + overrides: Final = {**environment, "DATABASE_URL": database_url} + with ( + owned_proxy_process( + shared, directory / "first", overrides, config=config, remove_environment=removed + ) as first, + owned_proxy_process( + shared, + directory / "second", + overrides, + config=config, + remove_environment=removed, + workers=2, + healthcheck=WorkerHealthcheck.LENIENT, + ) as second, + replace(first.gateway, database_url=database_url).scenario() as scenario, + ): + eventually(second.worker_pids, lambda pids: len(pids) == 2, seconds=120) + yield Pods(first, second, scenario, sinks) + + +@pytest.fixture(scope="module") +def provider() -> Iterator[Wire]: + with wire_server(_upstream) as wire: + yield wire + + +@pytest.fixture(scope="module") +def pods(provider: Wire, audit_sinks: SpanSinks, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Pods]: + with _two_pods(provider, audit_sinks, tmp_path_factory.mktemp("team-sampling")) as started: + yield started diff --git a/tests/integration/observability/test_passthrough_guardrail_scan_text.py b/tests/integration/observability/test_passthrough_guardrail_scan_text.py new file mode 100644 index 00000000000..7437f696d4c --- /dev/null +++ b/tests/integration/observability/test_passthrough_guardrail_scan_text.py @@ -0,0 +1,153 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.anthropic_sse import message_json +from integration._support.client import Gateway +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +GENERIC: Final = "passthrough-scan-generic" +AIRS: Final = "passthrough-scan-airs" +OPENAI_ROUTE: Final = "/openai/chat/completions" +ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" +ROUTE_BODIES: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + OPENAI_ROUTE: {"model": "gpt-5.6", "messages": [{"role": "user", "content": "Say hi in one word"}]}, + ANTHROPIC_ROUTE: { + "model": "claude-sonnet-5-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "Say hi in one word"}], + }, + } +) +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) + + +def _generic_verdict(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _airs_verdict(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "allow", "category": "benign", "scan_id": "s", "report_id": "r"}).encode()) + + +def _openai_reply(_request: Request) -> Reply: + return chat_reply("chatcmpl-passthrough-scan", "gpt-5.6", "Hi", stream=False) + + +def _anthropic_reply(_request: Request) -> Reply: + return Reply(body=message_json("msg_passthrough_scan", "claude-sonnet-5-5", "Hi")) + + +def _parsed_texts(texts: JsonValue, vendor_body: dict[str, JsonValue]) -> list[JsonValue]: + assert isinstance(texts, list), vendor_body + return [json.loads(text) if isinstance(text, str) else text for text in texts] + + +def _scanned_payloads(guardrail: str, vendor_body: dict[str, JsonValue]) -> list[JsonValue]: + if guardrail == GENERIC: + return _parsed_texts(vendor_body["texts"], vendor_body) + contents: Final = vendor_body["contents"] + assert isinstance(contents, list), vendor_body + return _parsed_texts([content.get("prompt") for content in contents if isinstance(content, dict)], vendor_body) + + +@dataclass(frozen=True, slots=True) +class PassThroughRig: + owned: OwnedProxy + vendors: Mapping[str, Wire] + upstreams: Mapping[str, Wire] + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[PassThroughRig]: + directory: Final = tmp_path_factory.mktemp("passthrough-guardrail-scan") + with ExitStack() as stack: + vendors: Final = MappingProxyType( + { + GENERIC: stack.enter_context(wire_server(_generic_verdict)), + AIRS: stack.enter_context(wire_server(_airs_verdict)), + } + ) + upstreams: Final = MappingProxyType( + { + OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), + ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)), + } + ) + config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + path: Final = directory / "passthrough-guardrail-scan.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "guardrails": [ + { + "guardrail_name": GENERIC, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": vendors[GENERIC].url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": AIRS, + "litellm_params": { + "guardrail": "panw_prisma_airs", + "mode": "pre_call", + "default_on": False, + "api_base": vendors[AIRS].url, + "api_key": "synthetic-airs-key", + "profile_name": "integration", + }, + }, + ], + } + ) + ) + environment: Final = { + "OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url, + "OPENAI_API_KEY": "synthetic-openai-key", + "ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + } + owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) + yield PassThroughRig(owned, vendors, upstreams) + + +@pytest.mark.parametrize("guardrail", [GENERIC, AIRS]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE]) +def test_pass_through_guardrail_scans_the_client_body_without_the_callers_key_metadata( + rig: PassThroughRig, route: str, guardrail: str +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + team: Final = scenario.team() + user: Final = scenario.user( + user_email=f"passthrough-scan-{uuid.uuid4().hex}@example.com", user_role="internal_user", teams=[team] + ) + key: Final = scenario.key(user_id=user, team_id=team) + body: Final = {**ROUTE_BODIES[route], "guardrails": [guardrail]} + + response: Final = gateway.request("POST", route, body, key=key) + + assert response.status_code == 200, response.text + scanned: Final = [ + _scanned_payloads(guardrail, _JSON.validate_json(request.body)) + for request in rig.vendors[guardrail].drain() + ] + assert scanned == [[body]], scanned + assert len(rig.upstreams[route].drain()) == 1 diff --git a/tests/integration/pricing/test_opencode_go_reported_cost_wire.py b/tests/integration/pricing/test_opencode_go_reported_cost_wire.py new file mode 100644 index 00000000000..6ac3bdccf27 --- /dev/null +++ b/tests/integration/pricing/test_opencode_go_reported_cost_wire.py @@ -0,0 +1,198 @@ +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_LIST_COST: Final = 0.00412 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + +# OpenCode's documented usage block (opencode.ai/v2/docs/console/api/inference, "Cost", read 2026-10-24): +# usage.cost is what hit the prepaid balance (0 inside the Go plan), usage.opencode.list_cost is the model price +_GO_BLOCK: Final[Mapping[str, JsonValue]] = { + "cost": 0, + "opencode": {"billing_source": "go", "list_cost": _LIST_COST, "allowance": _LIST_COST}, +} +_CHAT_USAGE: Final[Mapping[str, JsonValue]] = { + "prompt_tokens": 1200, + "completion_tokens": 300, + "total_tokens": 1500, + **_GO_BLOCK, +} +_RESPONSES_USAGE: Final[Mapping[str, JsonValue]] = { + "input_tokens": 1200, + "output_tokens": 300, + "total_tokens": 1500, + **_GO_BLOCK, +} +_MESSAGES_USAGE: Final[Mapping[str, JsonValue]] = {"input_tokens": 1200, "output_tokens": 300, **_GO_BLOCK} + + +def _approx(value: float) -> object: + return pytest.approx(value, rel=1e-6) # pyright: ignore[reportUnknownMemberType] # pytest lacks typed approx stubs + + +def _data_frame(event: Mapping[str, JsonValue]) -> bytes: + return f"data: {json.dumps(event)}\n\n".encode() + + +def _event_frame(event: Mapping[str, JsonValue]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _chat_completion() -> bytes: + return json.dumps( + { + "id": "chatcmpl-go", + "object": "chat.completion", + "created": 1, + "model": "glm-5.3", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "ok"}}], + "usage": dict(_CHAT_USAGE), + } + ).encode() + + +def _chat_stream() -> tuple[bytes, ...]: + text: Final = { + "id": "chatcmpl-go", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-5.3", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}], + } + finish: Final = {**text, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]} + final_usage: Final = {**text, "choices": [], "usage": dict(_CHAT_USAGE)} + return (*(_data_frame(event) for event in (text, finish, final_usage)), b"data: [DONE]\n\n") + + +def _responses_response() -> Mapping[str, JsonValue]: + return { + "id": "resp_go", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-6-luna", + "output": [ + { + "id": "msg_go", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": dict(_RESPONSES_USAGE), + } + + +def _responses_stream() -> tuple[bytes, ...]: + response: Final = _responses_response() + return ( + _event_frame( + {"type": "response.created", "sequence_number": 0, "response": {**response, "status": "in_progress"}} + ), + _event_frame( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_go", + "output_index": 0, + "content_index": 0, + "delta": "ok", + } + ), + _event_frame({"type": "response.completed", "sequence_number": 2, "response": response}), + ) + + +def _anthropic_message() -> Mapping[str, JsonValue]: + return { + "id": "msg_go", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": dict(_MESSAGES_USAGE), + } + + +def _anthropic_stream() -> tuple[bytes, ...]: + start: Final = {**_anthropic_message(), "content": [], "stop_reason": None} + return ( + _event_frame( + {"type": "message_start", "message": {**start, "usage": {"input_tokens": 1200, "output_tokens": 0}}} + ), + _event_frame({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _event_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}), + _event_frame({"type": "content_block_stop", "index": 0}), + _event_frame( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": dict(_MESSAGES_USAGE), + } + ), + _event_frame({"type": "message_stop"}), + ) + + +def _respond(request: Request) -> Reply: + assert request.method == "POST", request.method + streaming: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + if request.target.endswith("/chat/completions"): + return ( + Reply(content_type="text/event-stream", chunks=_chat_stream()) + if streaming + else Reply(body=_chat_completion()) + ) + if request.target.endswith("/responses"): + if streaming: + return Reply(content_type="text/event-stream", chunks=_responses_stream()) + return Reply(body=json.dumps(_responses_response()).encode()) + assert request.target.endswith("/v1/messages"), request.target + if streaming: + return Reply(content_type="text/event-stream", chunks=_anthropic_stream()) + return Reply(body=json.dumps(_anthropic_message()).encode()) + + +def _logged_spend(call_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + return float(rows[0]["spend"]) + + +def _body(path: str, model: str, stream: bool) -> Mapping[str, JsonValue]: + if path == "/v1/responses": + return {"model": model, "input": "hi", "stream": stream} + return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}], "stream": stream} + + +@pytest.mark.parametrize("stream", (False, True), ids=("non-stream", "stream")) +@pytest.mark.parametrize( + ("path", "model"), + ( + ("/v1/chat/completions", "opencode_go/glm-5.3"), + ("/v1/responses", "opencode_go/gpt-6-luna"), + ("/v1/messages", "opencode_go/claude-haiku-5-5"), + ), + ids=("chat", "responses", "messages"), +) +def test_request_is_billed_at_the_reported_list_cost(gateway: Gateway, path: str, model: str, stream: bool) -> None: + with wire_server(_respond) as wire, gateway.scenario() as scenario: + deployment: Final = scenario.model(model=model, api_base=wire.url, api_key="synthetic-go-key") + response: Final = gateway.request("POST", path, _body(path, deployment, stream)) + assert response.status_code == 200, response.text + assert "ok" in response.text + if not stream: + assert float(response.headers["x-litellm-response-cost"]) == _approx(_LIST_COST) + assert _logged_spend(response.headers["x-litellm-call-id"]) == _approx(_LIST_COST) diff --git a/tests/integration/providers/_api_base_url_join_support.py b/tests/integration/providers/_api_base_url_join_support.py new file mode 100644 index 00000000000..44bd6ec4dca --- /dev/null +++ b/tests/integration/providers/_api_base_url_join_support.py @@ -0,0 +1,163 @@ +import json +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +from integration._support.anthropic_sse import message_stream +from integration._support.openai_wire import chat_reply +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request +from integration.providers._api_base_trailing_slash_support import VECTOR, marker_of, reply_text, request_json +from pydantic import JsonValue + +COUNTED_TOKENS: Final = 7 +RELEVANCE: Final = 0.75 +_USAGE: Final = {"input_tokens": 3, "output_tokens": 2} + + +def _json(payload: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps(payload).encode()) + + +def _path(request: Request) -> str: + return urlsplit(request.target).path + + +def _wants_stream(request: Request) -> bool: + return request_json(request).get("stream") is True + + +def anthropic_message(marker: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": "scripted", + "content": [{"type": "text", "text": reply_text(marker)}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": dict(_USAGE), + } + + +def _anthropic(request: Request) -> Reply: + if _path(request).endswith("/count_tokens"): + return _json({"input_tokens": COUNTED_TOKENS}) + marker: Final = marker_of(request) + if _wants_stream(request): + return Reply( + content_type="text/event-stream", chunks=message_stream(f"msg_{marker}", "scripted", reply_text(marker)) + ) + return _json(anthropic_message(marker)) + + +def _gemini_candidate(text: str, finished: bool) -> dict[str, JsonValue]: + finish: Final[dict[str, JsonValue]] = {"finishReason": "STOP"} if finished else {} + return {"content": {"role": "model", "parts": [{"text": text}]}, "index": 0, **finish} + + +def _gemini_frame(text: str, finished: bool) -> bytes: + usage: Final = {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5} + frame: Final = {"candidates": [_gemini_candidate(text, finished)], "usageMetadata": usage} + return b"data: " + json.dumps(frame).encode() + b"\r\n\r\n" + + +def _gemini(request: Request) -> Reply: + path: Final = _path(request) + marker: Final = marker_of(request) + if path.endswith(":batchEmbedContents"): + return _json({"embeddings": [{"values": list(VECTOR)}]}) + if path.endswith(":streamGenerateContent"): + return Reply( + content_type="text/event-stream", + chunks=(_gemini_frame("reply ", False), _gemini_frame(marker, True)), + ) + return _json( + { + "candidates": [_gemini_candidate(reply_text(marker), True)], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, + } + ) + + +def _embedding_list(model: str) -> Reply: + return _json( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": list(VECTOR)}], + "model": model, + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + } + ) + + +def _openai_like(model: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + path: Final = _path(request) + marker: Final = marker_of(request) + if path.endswith("/embeddings"): + return _embedding_list(model) + if path.endswith("/audio/transcriptions"): + return _json({"text": reply_text(marker), "model": model}) + return chat_reply(f"chatcmpl-{marker}", model, reply_text(marker), stream=_wants_stream(request)) + + return respond + + +def _converse_event(kind: str, payload: Mapping[str, JsonValue]) -> bytes: + headers: Final = {":event-type": kind, ":content-type": "application/json", ":message-type": "event"} + return aws_event_stream_frame(headers, json.dumps(payload, separators=(",", ":")).encode()) + + +def _bedrock_converse(request: Request) -> Reply: + marker: Final = marker_of(request) + usage: Final[dict[str, JsonValue]] = {"inputTokens": 3, "outputTokens": 2, "totalTokens": 5} + if _path(request).endswith("/converse-stream"): + events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = ( + ("messageStart", {"role": "assistant"}), + ("contentBlockDelta", {"delta": {"text": "reply "}, "contentBlockIndex": 0}), + ("contentBlockDelta", {"delta": {"text": marker}, "contentBlockIndex": 0}), + ("contentBlockStop", {"contentBlockIndex": 0}), + ("messageStop", {"stopReason": "end_turn"}), + ("metadata", {"usage": usage, "metrics": {"latencyMs": 1}}), + ) + return Reply( + content_type="application/vnd.amazon.eventstream", + chunks=tuple(_converse_event(kind, payload) for kind, payload in events), + ) + return _json( + { + "output": {"message": {"role": "assistant", "content": [{"text": reply_text(marker)}]}}, + "stopReason": "end_turn", + "usage": usage, + "metrics": {"latencyMs": 1}, + } + ) + + +def _cohere_rerank(request: Request) -> Reply: + return _json( + { + "id": f"rerank-{marker_of(request)}", + "results": [{"index": 0, "relevance_score": RELEVANCE}], + "meta": {"billed_units": {"search_units": 1}}, + } + ) + + +_BY_PREFIX: Final[Mapping[str, Callable[[Request], Reply]]] = MappingProxyType( + { + "an": _anthropic, + "bi": lambda request: _json(anthropic_message(marker_of(request))), + "bc": _bedrock_converse, + "gm": _gemini, + "mi": _openai_like("mistral-small-latest"), + "db": _openai_like("databricks-meta-llama-3-3-70b-instruct"), + "co": _cohere_rerank, + } +) + + +def join_peer(request: Request) -> Reply: + return _BY_PREFIX[request.target.split("/")[1]](request) diff --git a/tests/integration/providers/test_api_base_url_join_chaos.py b/tests/integration/providers/test_api_base_url_join_chaos.py new file mode 100644 index 00000000000..8ebcd165f4d --- /dev/null +++ b/tests/integration/providers/test_api_base_url_join_chaos.py @@ -0,0 +1,176 @@ +import asyncio +from collections.abc import Awaitable, Callable +from itertools import chain, repeat +from typing import Final, TypeAlias + +import anthropic +import httpx +import openai +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.providers._api_base_trailing_slash_support import ( + KEY, + VECTOR, + assert_control_answers, + assert_one_request, + body_of, + marker_of, + new_marker, + of_marker, + reply_text, +) +from integration.providers._api_base_url_join_support import join_peer +from pydantic import JsonValue + +_PER_KIND: Final = 3 +_HAIKU: Final = "anthropic.claude-haiku-5-5" +_GENERATE: Final = "/gm/v1beta/models/gemini-3.8-flash:generateContent" +_STREAM_GENERATE: Final = "/gm/v1beta/models/gemini-3.8-flash:streamGenerateContent?alt=sse" + +Answer: TypeAlias = tuple[str | None, str | None] +Call: TypeAlias = Callable[[str], Awaitable[Answer]] + + +def _deploy(scenario: Scenario, model: str, api_base: str) -> str: + if model.startswith("bedrock/"): + return scenario.model(model=model, api_base=api_base, api_key=KEY, aws_region_name="us-east-1") + return scenario.model(model=model, api_base=api_base, api_key=KEY) + + +def _spend_api_bases(identities: tuple[str, ...]) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + "SELECT request_id, api_base FROM \"LiteLLM_SpendLogs\" WHERE request_id = ANY(string_to_array(%s, ','))", + (",".join(identities),), + ), + lambda found: len(found) == len(identities), + seconds=70, + ) + return {str(row["request_id"]): row["api_base"] for row in rows} + + +def test_mixed_burst_on_slash_bases_lands_each_call_once_on_its_single_slash_url(gateway: Gateway) -> None: + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + claude: Final = _deploy(scenario, "anthropic/claude-haiku-5-5", f"{wire.url}/an/") + gemini: Final = _deploy(scenario, "gemini/gemini-3.8-flash", f"{wire.url}/gm/v1beta/") + mistral: Final = _deploy(scenario, "mistral/mistral-small-latest", f"{wire.url}/mi/v1/") + converse: Final = _deploy(scenario, f"bedrock/converse/{_HAIKU}", f"{wire.url}/bc/") + databricks: Final = _deploy( + scenario, "databricks/databricks-meta-llama-3-3-70b-instruct", f"{wire.url}/db/serving-endpoints/" + ) + embed_model: Final = _deploy(scenario, "mistral/mistral-embed", f"{wire.url}/mi/v1/") + + def openai_client() -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + def chat(model: str) -> Call: + async def call(marker: str) -> Answer: + async with openai_client() as client: + completion: Final = await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}] + ) + return completion.choices[0].message.content, completion.id + + return call + + def streamed(model: str) -> Call: + async def call(marker: str) -> Answer: + async with openai_client() as client: + stream: Final = await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + parts: Final = [chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices] + return "".join(parts), None + + return call + + async def messages(marker: str) -> Answer: + async with anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0 + ) as client: + message: Final = await client.messages.create( + model=claude, max_tokens=16, messages=[{"role": "user", "content": marker}] + ) + return "".join(block.text for block in message.content if block.type == "text"), None + + async def embed(marker: str) -> Answer: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + response: Final = await client.post( + "/v1/embeddings", + json={"model": embed_model, "input": marker}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + data: Final = body_of(response)["data"] + vector_matches: Final = data == [{"object": "embedding", "index": 0, "embedding": list(VECTOR)}] + return (reply_text(marker) if vector_matches else str(data)), None + + kinds: Final[tuple[tuple[str, Call], ...]] = ( + ("/an/v1/messages", chat(claude)), + ("/an/v1/messages", streamed(claude)), + ("/an/v1/messages", messages), + (_GENERATE, chat(gemini)), + (_STREAM_GENERATE, streamed(gemini)), + ("/mi/v1/chat/completions", chat(mistral)), + ("/mi/v1/chat/completions", streamed(mistral)), + (f"/bc/model/{_HAIKU}/converse", chat(converse)), + (f"/bc/model/{_HAIKU}/converse-stream", streamed(converse)), + ("/db/serving-endpoints/chat/completions", chat(databricks)), + ("/mi/v1/embeddings", embed), + ) + repeated_kinds: Final = tuple(chain.from_iterable(repeat(kind, _PER_KIND) for kind in kinds)) + calls: Final = tuple((target, call, new_marker()) for target, call in repeated_kinds) + + async def burst() -> list[Answer]: + return await asyncio.gather(*(call(marker) for _, call, marker in calls)) + + answers: Final = asyncio.run(burst()) + assert [text for text, _ in answers] == [reply_text(marker) for *_, marker in calls] + received: Final = wire.drain() + assert len(received) == len(calls), sorted(request.target for request in received) + for target, _, marker in calls: + assert_one_request(of_marker(received, marker), target, marker) + logged: Final = { + identity: f"{wire.url}{target}" + for (target, _, _), (_, identity) in zip(calls, answers, strict=True) + if identity is not None + } + assert _spend_api_bases(tuple(logged)) == logged + + +def test_upstream_dropping_doomed_connections_mid_burst_fails_only_those_calls(gateway: Gateway) -> None: + markers: Final = tuple(new_marker() for _ in range(16)) + doomed: Final = frozenset(markers[::2]) + + def respond(request: Request) -> Reply: + return Reply(drop_connection=True) if marker_of(request) in doomed else join_peer(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + claude: Final = _deploy(scenario, "anthropic/claude-haiku-5-5", f"{wire.url}/an/") + gemini: Final = _deploy(scenario, "gemini/gemini-3.8-flash", f"{wire.url}/gm/v1beta/") + routed: Final = { + marker: (claude, "/an/v1/messages") if index % 4 < 2 else (gemini, _GENERATE) + for index, marker in enumerate(markers) + } + + async def call(marker: str) -> tuple[int, str]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + response: Final = await client.post( + "/v1/chat/completions", + json={"model": routed[marker][0], "messages": [{"role": "user", "content": marker}]}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + return response.status_code, response.text + + async def burst() -> list[tuple[int, str]]: + return await asyncio.gather(*(call(marker) for marker in markers)) + + answers: Final = dict(zip(markers, asyncio.run(burst()), strict=True)) + received: Final = wire.drain() + assert_control_answers(gateway, scenario) + for marker in doomed: + assert answers[marker][0] >= 500 and "error" in answers[marker][1], answers[marker] + assert {request.target for request in of_marker(received, marker)} == {routed[marker][1]}, marker + for marker in frozenset(markers) - doomed: + assert answers[marker][0] == 200 and reply_text(marker) in answers[marker][1], answers[marker] + assert_one_request(of_marker(received, marker), routed[marker][1], marker) diff --git a/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py b/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py new file mode 100644 index 00000000000..4a075138305 --- /dev/null +++ b/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py @@ -0,0 +1,140 @@ +import os +from collections.abc import Callable, Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment +from integration._support.forward_proxy import ForwardProxy, refusing_forward_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration.providers._api_base_trailing_slash_support import KEY, new_marker +from pydantic import JsonValue + +_AZURE_HOST: Final = "anthropic-prod.services.ai.azure.com" +_EMBEDDINGS: Final = "host-embeddings" +_RERANK: Final = "host-rerank" +_AZURE_CLAUDE: Final = "host-azure-claude" +_AZURE_COUNT: Final = "host-azure-count" +_EMPTY_BASE: Final = "host-empty-base" +_SCRUBBED_PREFIXES: Final = ("ANTHROPIC_", "AZURE_AI_", "INFINITY_") +_PROXY_VARIABLES: Final = frozenset({"HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY"}) +_load_yaml: Final[Callable[[str], object]] = yaml.safe_load + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + + +def _deployment(name: str, model: str, api_base: str) -> dict[str, JsonValue]: + return {"model_name": name, "litellm_params": {"model": model, "api_base": api_base, "api_key": KEY}} + + +def _config(directory: Path) -> Path: + base: Final = JSON_OBJECT.validate_python(_load_yaml(Path("tests/integration/proxy_config.yaml").read_text())) + configuration: Final = { + **base, + "model_list": [ + _deployment(_EMBEDDINGS, "infinity/BAAI/bge-small-en-v1.5", "http://embeddings"), + _deployment(_RERANK, "infinity/BAAI/bge-reranker-base", "http://rerank"), + _deployment(_AZURE_CLAUDE, "azure_ai/claude-haiku-5-5", f"http://{_AZURE_HOST}/anthropic"), + _deployment(_AZURE_COUNT, "azure_ai/claude-haiku-5-5", f"http://{_AZURE_HOST}"), + _deployment(_EMPTY_BASE, "anthropic/claude-haiku-5-5", ""), + ], + "router_settings": {"disable_cooldowns": True, "num_retries": 0}, + } + path: Final = directory / "host-named-like-the-route.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +def _scrubbed() -> tuple[str, ...]: + return tuple( + name + for name in os.environ + if name.upper() in _PROXY_VARIABLES or name.upper().startswith(_SCRUBBED_PREFIXES) + ) + + +@pytest.fixture(scope="module") +def forward() -> Iterator[ForwardProxy]: + with refusing_forward_proxy() as proxy: + yield proxy + + +@pytest.fixture(scope="module") +def isolated(forward: ForwardProxy, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("url-join-host") + with gateway_from_environment() as shared, owned_proxy( + shared, + directory, + { + "HTTP_PROXY": forward.url, + "HTTPS_PROXY": forward.url, + "NO_PROXY": "127.0.0.1,localhost", + "PROXY_BATCH_POLLING_ENABLED": "false", + }, + config=_config(directory), + remove_environment=_scrubbed(), + ) as candidate: + yield candidate + + +def _targets_after(forward: ForwardProxy, call: Callable[[], object]) -> frozenset[str]: + forward.drain() + call() + targets: Final = frozenset(forward.targets()) + assert targets, "the forward proxy saw no outbound request" + return targets + + +def _chat(isolated: Gateway, model: str) -> object: + return isolated.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": new_marker()}]} + ) + + +def _messages(isolated: Gateway, model: str, path: str) -> object: + return isolated.request( + "POST", + path, + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": new_marker()}]}, + headers={"anthropic-version": "2023-06-01"}, + ) + + +def test_embedding_host_named_embeddings_gets_the_route_appended(isolated: Gateway, forward: ForwardProxy) -> None: + targets: Final = _targets_after( + forward, lambda: isolated.request("POST", "/v1/embeddings", {"model": _EMBEDDINGS, "input": new_marker()}) + ) + assert targets == {"http://embeddings/embeddings"} + + +def test_rerank_host_named_rerank_gets_the_route_appended(isolated: Gateway, forward: ForwardProxy) -> None: + body: Final[dict[str, JsonValue]] = {"model": _RERANK, "query": new_marker(), "documents": ["only document"]} + targets: Final = _targets_after(forward, lambda: isolated.request("POST", "/v1/rerank", body)) + assert targets == {"http://rerank/rerank"} + + +def test_azure_claude_chat_on_a_host_starting_with_anthropic_keeps_the_host( + isolated: Gateway, forward: ForwardProxy +) -> None: + targets: Final = _targets_after(forward, lambda: _chat(isolated, _AZURE_CLAUDE)) + assert targets == {f"http://{_AZURE_HOST}/anthropic/v1/messages"} + + +def test_azure_claude_messages_on_a_host_starting_with_anthropic_keeps_the_host( + isolated: Gateway, forward: ForwardProxy +) -> None: + targets: Final = _targets_after(forward, lambda: _messages(isolated, _AZURE_CLAUDE, "/v1/messages")) + assert targets == {f"http://{_AZURE_HOST}/anthropic/v1/messages"} + + +def test_azure_claude_count_tokens_on_a_bare_host_starting_with_anthropic_adds_the_anthropic_path( + isolated: Gateway, forward: ForwardProxy +) -> None: + targets: Final = _targets_after(forward, lambda: _messages(isolated, _AZURE_COUNT, "/v1/messages/count_tokens")) + assert targets == {f"http://{_AZURE_HOST}/anthropic/v1/messages/count_tokens"} + + +def test_empty_api_base_falls_back_to_the_provider_default_host(isolated: Gateway, forward: ForwardProxy) -> None: + targets: Final = _targets_after(forward, lambda: _chat(isolated, _EMPTY_BASE)) + assert targets == {"api.anthropic.com:443"} diff --git a/tests/integration/providers/test_api_base_url_join_wire.py b/tests/integration/providers/test_api_base_url_join_wire.py new file mode 100644 index 00000000000..42e53ff9f95 --- /dev/null +++ b/tests/integration/providers/test_api_base_url_join_wire.py @@ -0,0 +1,373 @@ +import json +from typing import Final +from urllib.parse import urlsplit + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.providers._api_base_trailing_slash_support import ( + KEY, + TRAILING, + TRAILING_IDS, + VECTOR, + Row, + assert_control_answers, + assert_one_request, + body_of, + first_item, + new_marker, + reply_text, + request_json, +) +from integration.providers._api_base_url_join_support import COUNTED_TOKENS, RELEVANCE, join_peer +from pydantic import JsonValue + +_HAIKU: Final = "anthropic.claude-haiku-5-5" +_ANTHROPIC: Final = Row("anthropic/claude-haiku-5-5", "/an", "/v1/messages") +_GEMINI: Final = Row("gemini/gemini-3.8-flash", "/gm/v1beta", "/models/gemini-3.8-flash:generateContent") +_MISTRAL: Final = Row("mistral/mistral-small-latest", "/mi/v1", "/chat/completions") +_BEDROCK_INVOKE: Final = Row(f"bedrock/invoke/{_HAIKU}", "/bi", f"/model/{_HAIKU}/invoke") +_CHAT_ROWS: Final = ( + _ANTHROPIC, + _GEMINI, + _MISTRAL, + Row("databricks/databricks-meta-llama-3-3-70b-instruct", "/db/serving-endpoints", "/chat/completions"), + _BEDROCK_INVOKE, + Row(f"bedrock/converse/{_HAIKU}", "/bc", f"/model/{_HAIKU}/converse"), +) +_STREAM_ROWS: Final = ( + _ANTHROPIC, + Row("gemini/gemini-3.8-flash", "/gm/v1beta", "/models/gemini-3.8-flash:streamGenerateContent?alt=sse"), + _MISTRAL, + Row(f"bedrock/converse/{_HAIKU}", "/bc", f"/model/{_HAIKU}/converse-stream"), +) +_MESSAGES_ROWS: Final = (_ANTHROPIC, _BEDROCK_INVOKE, _GEMINI, _MISTRAL) +_RESPONSES_ROWS: Final = (_ANTHROPIC, _GEMINI) +_EMBEDDING_ROWS: Final = ( + Row("mistral/mistral-embed", "/mi/v1", "/embeddings"), + Row("gemini/gemini-embedding-2", "/gm/v1beta", "/models/gemini-embedding-2:batchEmbedContents"), +) +_TRANSCRIPTION: Final = Row("mistral/voxtral-mini-latest", "/mi/v1", "/audio/transcriptions") +_FULL_ROUTE_SUFFIXES: Final = ("", "/", "?beta=true") +_FULL_ROUTE_IDS: Final = ("exact", "slash", "query") + + +def _ids(rows: tuple[Row, ...]) -> tuple[str, ...]: + return tuple(row.model for row in rows) + + +def _deploy(scenario: Scenario, row: Row, api_base: str) -> str: + if row.model.startswith("bedrock/"): + return scenario.model(model=row.model, api_base=api_base, api_key=KEY, aws_region_name="us-east-1") + return scenario.model(model=row.model, api_base=api_base, api_key=KEY) + + +def _openai(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _chat(gateway: Gateway, model: str, marker: str) -> str | None: + completion: Final = _openai(gateway).chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}] + ) + return completion.choices[0].message.content + + +def _message_text(gateway: Gateway, model: str, marker: str) -> list[str]: + message: Final = _anthropic(gateway).messages.create( + model=model, max_tokens=16, messages=[{"role": "user", "content": marker}] + ) + return [block.text for block in message.content if block.type == "text"] + + +def _spend_rows(pattern: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, api_base, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', (pattern,) + ) + + +def _one_spend_row(pattern: str) -> dict[str, JsonValue]: + (row,) = eventually(lambda: _spend_rows(pattern), lambda rows: len(rows) == 1, seconds=70) + return row + + +def _gemini_contents(marker: str) -> dict[str, JsonValue]: + return {"contents": [{"role": "user", "parts": [{"text": marker}]}]} + + +def _gemini_text(payload: dict[str, JsonValue]) -> JsonValue: + content: Final = object_value(first_item(payload["candidates"])["content"]) + return first_item(content["parts"])["text"] + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _CHAT_ROWS, ids=_ids(_CHAT_ROWS)) +def test_chat_completion_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + assert _chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _STREAM_ROWS, ids=_ids(_STREAM_ROWS)) +def test_streamed_chat_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + stream: Final = _openai(gateway).chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + text: Final = "".join(chunk.choices[0].delta.content or "" for chunk in stream if chunk.choices) + assert text == reply_text(marker) + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _MESSAGES_ROWS, ids=_ids(_MESSAGES_ROWS)) +def test_messages_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + assert _message_text(gateway, model, marker) == [reply_text(marker)] + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_streamed_messages_reach_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an{trailing}") + with _anthropic(gateway).messages.stream( + model=model, max_tokens=16, messages=[{"role": "user", "content": marker}] + ) as stream: + text: Final = "".join(stream.text_stream) + assert text == reply_text(marker) + assert request_json(assert_one_request(wire.drain(), "/an/v1/messages", marker))["stream"] is True + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _RESPONSES_ROWS, ids=_ids(_RESPONSES_ROWS)) +def test_responses_bridge_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + assert _openai(gateway).responses.create(model=model, input=marker).output_text == reply_text(marker) + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _EMBEDDING_ROWS, ids=_ids(_EMBEDDING_ROWS)) +def test_embedding_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + response: Final = gateway.request("POST", "/v1/embeddings", {"model": model, "input": marker}) + assert first_item(body_of(response)["data"])["embedding"] == list(VECTOR) + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_mistral_transcription_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _TRANSCRIPTION, f"{wire.url}/mi/v1{trailing}") + response: Final = gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model}, + {"file": ("speech.wav", f"RIFF {marker}".encode(), "audio/wav")}, + ) + assert body_of(response)["text"] == reply_text(marker) + assert_one_request(wire.drain(), _TRANSCRIPTION.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_google_genai_generate_content_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _GEMINI, f"{wire.url}/gm/v1beta{trailing}") + response: Final = gateway.request( + "POST", f"/v1beta/models/{model}:generateContent", _gemini_contents(marker) + ) + assert _gemini_text(body_of(response)) == reply_text(marker) + assert_one_request(wire.drain(), _GEMINI.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_google_genai_stream_generate_content_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _GEMINI, f"{wire.url}/gm/v1beta{trailing}") + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:streamGenerateContent", + _gemini_contents(marker), + params={"alt": "sse"}, + ) + assert response.status_code == 200, response.text + assert marker in response.text, response.text + assert_one_request( + wire.drain(), "/gm/v1beta/models/gemini-3.8-flash:streamGenerateContent?alt=sse", marker + ) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_count_tokens_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an{trailing}") + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + assert body_of(response)["input_tokens"] == COUNTED_TOKENS + assert_one_request(wire.drain(), "/an/v1/messages/count_tokens", marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_health_check_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an{trailing}") + report: Final = body_of(gateway.request("GET", "/health", params={"model": model})) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/an/v1/messages")] + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_spend_log_records_the_single_slash_url_it_called(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an{trailing}") + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + identity: Final = string_value(body_of(response)["id"]) + assert_one_request(wire.drain(), "/an/v1/messages", marker) + assert _one_spend_row(identity)["api_base"] == f"{wire.url}/an/v1/messages" + + +def test_cache_hit_on_a_slash_base_skips_the_upstream_and_logs_its_own_row(gateway: Gateway) -> None: + pytest.skip("BUG: a response-cache hit writes two cache_hit spend rows for one chat completion") + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/") + body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "user", "content": marker}]} + first: Final = body_of(gateway.request("POST", "/v1/chat/completions", body)) + second: Final = body_of(gateway.request("POST", "/v1/chat/completions", body)) + assert first_item(second["choices"])["message"] == first_item(first["choices"])["message"] + assert_one_request(wire.drain(), "/an/v1/messages", marker) + identity: Final = string_value(first["id"]) + assert _one_spend_row(identity)["api_base"] == f"{wire.url}/an/v1/messages" + hit: Final = _one_spend_row(f"{identity}_cache_hit%") + assert hit["cache_hit"] == "True", hit + assert "//" not in urlsplit(string_value(hit["api_base"])).path, hit + + +@pytest.mark.parametrize("suffix", _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS) +def test_chat_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/v1/messages{suffix}") + assert _chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), f"/an/v1/messages{suffix.rstrip('/')}", marker) + + +@pytest.mark.parametrize("suffix", _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS) +def test_messages_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/v1/messages{suffix}") + assert _message_text(gateway, model, marker) == [reply_text(marker)] + assert_one_request(wire.drain(), f"/an/v1/messages{suffix.rstrip('/')}", marker) + + +@pytest.mark.parametrize("suffix", ("", "/", "?tier=a"), ids=("exact", "slash", "query")) +def test_rerank_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, Row("cohere/rerank-v3.5", "/co", ""), f"{wire.url}/co/v2/rerank{suffix}") + response: Final = gateway.request( + "POST", "/v1/rerank", {"model": model, "query": marker, "documents": ["only document"]} + ) + assert [object_value(result)["relevance_score"] for result in list_value(body_of(response)["results"])] == [ + RELEVANCE + ] + assert_one_request(wire.drain(), f"/co/v2/rerank{suffix.rstrip('/')}", marker) + + +@pytest.mark.parametrize( + ("row", "trailing"), + ( + (_ANTHROPIC, "//"), + (Row("anthropic/claude-haiku-5-5", f"/an/{'p' * 5000}", "/v1/messages"), "/"), + (_GEMINI, "//"), + (Row("gemini/gemini-3.8-flash", f"/gm/{'p' * 5000}", "/models/gemini-3.8-flash:generateContent"), "/"), + ), + ids=("anthropic-two-slashes", "anthropic-5kb-prefix", "gemini-two-slashes", "gemini-5kb-prefix"), +) +def test_base_with_repeated_slashes_or_a_long_prefix_reaches_one_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + assert _chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), row.target, marker) + + +def _anthropic_error(status: int, marker: str) -> bytes: + return json.dumps( + {"type": "error", "error": {"type": "authentication_error", "message": f"scripted {status} {marker}"}} + ).encode() + + +def _gemini_error(status: int, marker: str) -> bytes: + error: Final = {"code": status, "message": f"scripted {status} {marker}", "status": "NOT_FOUND"} + return json.dumps({"error": error}).encode() + + +@pytest.mark.parametrize( + ("row", "status"), + ((_ANTHROPIC, 401), (_GEMINI, 404)), + ids=("anthropic-401", "gemini-404"), +) +def test_upstream_error_on_a_slash_base_reaches_the_caller_from_one_url( + gateway: Gateway, row: Row, status: int +) -> None: + marker: Final = new_marker() + error: Final = _anthropic_error(status, marker) if row is _ANTHROPIC else _gemini_error(status, marker) + + def respond(_: Request) -> Reply: + return Reply(status=status, body=error) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}/") + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + assert response.status_code == status, response.text + assert marker in response.text, response.text + assert_one_request(wire.drain(), row.target, marker) + assert_control_answers(gateway, scenario) + + +def test_streamed_chat_upstream_error_on_a_slash_base_reaches_the_caller(gateway: Gateway) -> None: + marker: Final = new_marker() + + def respond(_: Request) -> Reply: + return Reply(status=401, body=_anthropic_error(401, marker)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/") + with pytest.raises(openai.AuthenticationError) as raised: + _openai(gateway).chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + assert marker in str(raised.value) + assert_one_request(wire.drain(), "/an/v1/messages", marker) + assert_control_answers(gateway, scenario) diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index 2badf810b6a..d4f9533fcdc 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -1,6 +1,7 @@ import json import re from collections.abc import Callable, Mapping +from itertools import chain from typing import Final from uuid import uuid4 @@ -263,13 +264,17 @@ def _assert_chunk_envelope(chunks: tuple[dict[str, JsonValue], ...], model: str, return next(iter(identities)) +def _content_chunks(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(chunk for chunk in chunks if chunk.get("usage") is None) + + def _finish_reasons(chunks: tuple[dict[str, JsonValue], ...]) -> list[JsonValue]: - reasons: Final = (_only_choice(chunk).get("finish_reason") for chunk in chunks) + reasons: Final = (_only_choice(chunk).get("finish_reason") for chunk in _content_chunks(chunks)) return [reason for reason in reasons if reason is not None] def _deltas(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: - return tuple(object_value(_only_choice(chunk)["delta"]) for chunk in chunks) + return tuple(object_value(_only_choice(chunk)["delta"]) for chunk in _content_chunks(chunks)) def _sdk_base(gateway: Gateway) -> str: @@ -301,12 +306,12 @@ def test_chat_completion_on_a_mantle_claude_id_posts_native_messages_and_spends_ @pytest.mark.parametrize( ("stream_options", "usage_chunks"), [ - pytest.param({"stream_options": {"include_usage": True}}, [_USAGE], id="include_usage"), + pytest.param({"stream_options": {"include_usage": True}}, [(True, *_USAGE)], id="include_usage"), pytest.param({}, [], id="no_usage"), ], ) def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_stream( - gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[int, int, int]] + gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[JsonValue, ...]] ) -> None: marker: Final = uuid4().hex expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) @@ -322,9 +327,9 @@ def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_ identity: Final = _assert_chunk_envelope(chunks, model, response.text) assert "".join(str(delta.get("content") or "") for delta in _deltas(chunks)) == _answer(marker), response.text assert _finish_reasons(chunks) == ["stop"], response.text - assert [_token_usage(chunk) for chunk in chunks if chunk.get("usage") is not None] == usage_chunks, ( - response.text - ) + assert [ + (chunk["choices"] == [], *_token_usage(chunk)) for chunk in chunks if chunk.get("usage") is not None + ] == usage_chunks, response.text _assert_success_row( _spend_row(_BY_REQUEST_ID, identity), model, _HAIKU, "acompletion", _mantle_cost(gateway, _HAIKU) ) @@ -620,15 +625,14 @@ def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: ) identities: Final = {chunk.id for chunk in chunks} assert len(identities) == 1 and {chunk.model for chunk in chunks} == {model}, chunks - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _answer(marker), chunks - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"], ( - chunks - ) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _answer(marker), chunks + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"], chunks assert [ - (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + (len(chunk.choices), chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None - ] == [_USAGE], chunks + ] == [(0, *_USAGE)], chunks _assert_success_row( _spend_row(_BY_REQUEST_ID, next(iter(identities))), model, @@ -688,15 +692,14 @@ async def test_async_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gatew chunks: Final = tuple([chunk async for chunk in stream]) identities: Final = {chunk.id for chunk in chunks} assert len(identities) == 1 and {chunk.model for chunk in chunks} == {model}, chunks - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _answer(marker), chunks - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"], ( - chunks - ) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _answer(marker), chunks + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"], chunks assert [ - (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + (len(chunk.choices), chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None - ] == [_USAGE], chunks + ] == [(0, *_USAGE)], chunks _assert_success_row( _spend_row(_BY_REQUEST_ID, next(iter(identities))), model, diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index ea29914ba65..89ca674ef31 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -660,3 +660,22 @@ def test_a_deployment_whose_provider_has_no_decisions_support_is_refused_naming_ for provider in _PROVIDERS: assert provider.name in response.text, response.text assert _upstream_calls(gateway, handle) == [] + + +def test_typesafe_gateway_preserves_missing_usage(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + body = {"model": "jev-1.13.0", "answers": {"exp": {"type": "noul", "noul": 0.8}}} + handle = _register(scenario, body) + model = _deployment(scenario, handle, _PROVIDERS[1]) + response = gateway.request( + "POST", + "/v1/systemone", + { + "model": model, + "state": {"finding": "x"}, + "questions": {"exp": {"type": "noul", "instructions": "Is it exploitable?"}}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["usage"] is None + assert response.json()["answers"] == body["answers"] diff --git a/tests/integration/providers/test_responses_compaction_items_wire.py b/tests/integration/providers/test_responses_compaction_items_wire.py new file mode 100644 index 00000000000..5b8e2b534d7 --- /dev/null +++ b/tests/integration/providers/test_responses_compaction_items_wire.py @@ -0,0 +1,1057 @@ +import json +import os +import re +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +import httpx +import openai +import pytest +import yaml +from integration._support import claude_code as cc +from integration._support.client import ( + 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.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from pydantic import JsonValue, TypeAdapter + +from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD +from litellm.responses.utils import ResponsesAPIRequestUtils + +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_CLAUDE: Final = cc.OPUS +_GEMINI: Final = "gemini-3.8-flash" +_OPENAI: Final = "gpt-5.6" +_PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) +_CACHE_BUST: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}} +_COMPACTING: Final[list[JsonValue]] = [{"type": "compaction", "compact_threshold": 150000}] +_COMPACT_BETA: Final = "compact-2026-01-12" +_CITY: Final = "Tokyo" +_MASKED_CITY: Final = "[CITY]" +_MASK_INPUT: Final = "compaction-mask-input" +_MASK_OUTPUT: Final = "compaction-mask-output" +_BLOCK_OUTPUT: Final = "compaction-block-output" +_REDACT_HEADER: Final = {"x-litellm-enable-message-redaction": "true"} +_UNREDACT_HEADER: Final = {"litellm-disable-message-redaction": "true"} +_SCRIPTED_REJECTION: Final = json.dumps( + {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted rejection"}} +).encode() +_ASK: Final = re.compile(r"ask-([0-9a-f]{32})") + +Route: TypeAlias = Literal["claude-compacting", "claude-plain", "gemini", "openai", "openai-compact", "openai-opaque"] + + +def _new_marker() -> str: + return uuid.uuid4().hex + + +def _ask(marker: str, *, compact: bool = False) -> str: + return f"ask-{marker} compact-now-{marker}" if compact else f"ask-{marker}" + + +def _fail(marker: str) -> str: + return f"fail-now-{marker}" + + +def _think(marker: str) -> str: + return f"think-now-{marker}" + + +def _thinking(marker: str) -> dict[str, JsonValue]: + return {"type": "thinking", "thinking": f"thought-{marker}", "signature": f"signature-{marker}"} + + +def _answer(marker: str) -> str: + return f"answer-{marker}" + + +def _summary(marker: str) -> str: + return f"kept-{marker} {'a' * 1500} middle-{marker} {'b' * 1500} the user lives in {_CITY}" + + +def _block(summary: str) -> dict[str, JsonValue]: + return {"type": "compaction", "content": summary} + + +def _compaction_item(summary: str) -> dict[str, JsonValue]: + return { + "type": "compaction", + "id": f"cmp_{uuid.uuid4().hex}", + "encrypted_content": json.dumps(_block(summary), separators=(",", ":")), + } + + +def _newest_marker(raw: bytes) -> str: + found: Final = _ASK.findall(raw.decode()) + assert found, raw[:400] + return str(found[-1]) + + +def _anthropic_usage() -> dict[str, JsonValue]: + return {"input_tokens": 20, "output_tokens": 5} + + +def _anthropic_block_frames(index: int, block: Mapping[str, JsonValue]) -> tuple[bytes, ...]: + def start(content_block: Mapping[str, JsonValue]) -> bytes: + return cc.sse_frame( + "content_block_start", {"type": "content_block_start", "index": index, "content_block": dict(content_block)} + ) + + def delta(fields: Mapping[str, JsonValue]) -> bytes: + return cc.sse_frame( + "content_block_delta", {"type": "content_block_delta", "index": index, "delta": dict(fields)} + ) + + stop: Final = cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": index}) + match block["type"]: + case "compaction": + summary: Final = string_value(block["content"]) + half: Final = len(summary) // 2 + return ( + start({"type": "compaction", "content": ""}), + delta({"type": "compaction_delta", "content": summary[:half]}), + delta({"type": "compaction_delta", "content": summary[half:]}), + stop, + ) + case "thinking": + return ( + start({"type": "thinking", "thinking": ""}), + delta({"type": "thinking_delta", "thinking": block["thinking"]}), + delta({"type": "signature_delta", "signature": block["signature"]}), + stop, + ) + case _: + return (start({"type": "text", "text": ""}), delta({"type": "text_delta", "text": block["text"]}), stop) + + +def _anthropic_stream(identity: str, content: Sequence[Mapping[str, JsonValue]]) -> tuple[bytes, ...]: + start: Final = cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": _CLAUDE, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 1}, + }, + }, + ) + blocks: Final = tuple( + chain.from_iterable(_anthropic_block_frames(index, block) for index, block in enumerate(content)) + ) + end: Final = cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + return (start, *blocks, end, cc.sse_frame("message_stop", {"type": "message_stop"})) + + +def _claude_reply(body: Mapping[str, JsonValue], raw: bytes) -> Reply: + marker: Final = _newest_marker(raw) + if _fail(marker).encode() in raw: + return Reply(status=400, body=_SCRIPTED_REJECTION) + compaction: Final = (_block(_summary(marker)),) if f"compact-now-{marker}".encode() in raw else () + thinking: Final = (_thinking(marker),) if _think(marker).encode() in raw else () + content: Final = (*compaction, *thinking, {"type": "text", "text": _answer(marker)}) + identity: Final = f"msg_{uuid.uuid4().hex}" + if body.get("stream") is True: + return Reply(content_type="text/event-stream", chunks=_anthropic_stream(identity, content)) + return Reply(body=cc.message_reply(identity, _CLAUDE, content, _anthropic_usage())) + + +def _gemini_reply(raw: bytes) -> Reply: + frame: Final = { + "responseId": f"gemini-{uuid.uuid4().hex}", + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": _answer(_newest_marker(raw))}]}, + "index": 0, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 5, "totalTokenCount": 25}, + "modelVersion": _GEMINI, + } + return Reply(body=json.dumps(frame).encode()) + + +_OPENAI_USAGE: Final = { + "input_tokens": 20, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 5, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 25, +} + + +def _openai_reply(raw: bytes) -> Reply: + tag: Final = uuid.uuid4().hex + message: Final = { + "id": f"msg_{tag}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": _answer(_newest_marker(raw)), "annotations": []}], + } + response: Final = { + "id": f"resp_{tag}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": _OPENAI, + "output": [message], + "usage": _OPENAI_USAGE, + } + return Reply(body=json.dumps(response).encode()) + + +def _openai_compact_reply() -> Reply: + compacted: Final = { + "id": f"resp_compact_{uuid.uuid4().hex}", + "created_at": 1, + "object": "response.compaction", + "output": [], + "usage": _OPENAI_USAGE, + } + return Reply(body=json.dumps(compacted).encode()) + + +def _vendor(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": [{"id": _OPENAI, "object": "model"}]}).encode()) + path: Final = urlsplit(request.target).path + if path.endswith("/v1/messages/count_tokens"): + return Reply(body=json.dumps({"input_tokens": _word_count(_JSON_OBJECT.validate_json(request.body))}).encode()) + if path.endswith("/v1/messages"): + return _claude_reply(_JSON_OBJECT.validate_json(request.body), request.body) + if path.endswith(":generateContent"): + return _gemini_reply(request.body) + if path.endswith("/responses/compact"): + return _openai_compact_reply() + assert path.endswith("/responses"), request.target + return _openai_reply(request.body) + + +def _provider_calls(wire: Wire, marker: str) -> tuple[Request, ...]: + needle: Final = f"ask-{marker}".encode() + return tuple(request for request in wire.drain() if request.method == "POST" and needle in request.body) + + +def _only_call(wire: Wire, marker: str) -> tuple[Request, dict[str, JsonValue]]: + calls: Final = _provider_calls(wire, marker) + assert len(calls) == 1, [call.target for call in calls] + return calls[0], _JSON_OBJECT.validate_json(calls[0].body) + + +def _objects(value: JsonValue) -> tuple[dict[str, JsonValue], ...]: + assert isinstance(value, list), value + return tuple(object_value(item) for item in value) + + +def _message_texts(message: Mapping[str, JsonValue], key: str = "content") -> tuple[str, ...]: + content: Final = message.get(key) + if isinstance(content, str): + return (content,) + if not isinstance(content, list): + return () + return tuple(string_value(part["text"]) for part in _objects(content) if isinstance(part.get("text"), str)) + + +def _anthropic_texts(body: Mapping[str, JsonValue], role: str) -> tuple[str, ...]: + messages: Final = tuple(message for message in _objects(body["messages"]) if message["role"] == role) + return tuple(chain.from_iterable(map(_message_texts, messages))) + + +def _word_count(body: Mapping[str, JsonValue]) -> int: + texts: Final = (*_anthropic_texts(body, "user"), *_anthropic_texts(body, "assistant")) + return sum(len(text.split()) for text in texts) + + +def _anthropic_block_types(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + contents: Final = tuple(message["content"] for message in _objects(body["messages"])) + blocks: Final = chain.from_iterable(_objects(content) for content in contents if isinstance(content, list)) + return tuple(block.get("type") for block in blocks) + + +def _gemini_texts(body: Mapping[str, JsonValue], role: str) -> tuple[str, ...]: + turns: Final = tuple(turn for turn in _objects(body["contents"]) if turn["role"] == role) + return tuple(chain.from_iterable(_message_texts(turn, "parts") for turn in turns)) + + +def _openai_user_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]: + items: Final = tuple(item for item in _objects(body["input"]) if item.get("role") == "user") + return tuple(chain.from_iterable(map(_message_texts, items))) + + +def _openai_item_types(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + return tuple(item.get("type") for item in _objects(body["input"])) + + +def _register(scenario: Scenario, wire: Wire, provider: Literal["claude", "gemini", "openai"]) -> str: + match provider: + case "claude": + model = scenario.model(model=f"anthropic/{_CLAUDE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + case "gemini": + model = scenario.model(model=f"gemini/{_GEMINI}", api_base=wire.url, api_key="synthetic-gemini-key") + case "openai": + model = scenario.model(model=f"openai/{_OPENAI}", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key") + wait_for_worker_readiness(scenario.gateway, model, 1, _PROXY_WORKERS) + return model + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _decoded(encrypted_content: str) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(encrypted_content) + + +def _responses( + gateway: Gateway, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> dict[str, JsonValue]: + response: Final = gateway.request("POST", "/v1/responses", {**body, **_CACHE_BUST}, headers=headers) + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: + payloads: Final = tuple( + line.removeprefix("data:").strip() for line in text.splitlines() if line.startswith("data:") + ) + return tuple(_JSON_OBJECT.validate_json(payload) for payload in payloads if payload != "[DONE]") + + +def _stream( + gateway: Gateway, path: str, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> tuple[dict[str, JsonValue], ...]: + sent_headers: Final = {"Authorization": f"Bearer {gateway.key}", **(headers or {})} + with gateway.client.stream( + "POST", path, json={**body, **_CACHE_BUST, "stream": True}, headers=sent_headers + ) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + return _sse_events(text) + + +def _delta_provider_fields(chunk: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + fields + for choice in _objects(chunk.get("choices", [])) + if isinstance(fields := object_value(choice.get("delta", {})).get("provider_specific_fields"), dict) + ) + + +def _streamed_compaction_blocks(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + fields: Final = chain.from_iterable(_delta_provider_fields(chunk) for chunk in chunks) + return tuple(field["compaction_blocks"] for field in fields if "compaction_blocks" in field) + + +def _output_items(response: Mapping[str, JsonValue], item_type: str) -> tuple[dict[str, JsonValue], ...]: + return tuple(item for item in _objects(response["output"]) if item["type"] == item_type) + + +def _listed(value: JsonValue) -> list[JsonValue]: + assert isinstance(value, list), value + return value + + +def _only(items: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + assert len(items) == 1, items + return items[0] + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +def test_claude_compaction_block_is_returned_as_a_compaction_output_item(gateway: Gateway, stream: bool) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + client: Final = _openai_client(gateway) + request: Final = { + "input": _ask(marker, compact=True), + "extra_body": {"context_management": _COMPACTING, **_CACHE_BUST}, + } + if not stream: + reply: Final = client.responses.create(model=model, **request) + assert [item.type for item in reply.output] == ["compaction", "message"] + compaction: Final = reply.output[0] + assert compaction.type == "compaction" + assert _decoded(compaction.encrypted_content) == _block(_summary(marker)) + assert reply.output_text == _answer(marker) + else: + events: Final = tuple(client.responses.create(model=model, stream=True, **request)) + added: Final = [e for e in events if e.type == "response.output_item.added" and e.item.type == "compaction"] + done: Final = [e for e in events if e.type == "response.output_item.done" and e.item.type == "compaction"] + completed: Final = [e for e in events if e.type == "response.completed"] + assert len(added) == 1 and len(done) == 1 and len(completed) == 1, [e.type for e in events] + assert added[0].item.encrypted_content == "" + assert done[0].output_index == added[0].output_index + assert done[0].item.id == added[0].item.id + assert _decoded(done[0].item.encrypted_content) == _block(_summary(marker)) + finished: Final = [item for item in completed[0].response.output if item.type == "compaction"] + assert [(item.id, item.encrypted_content) for item in finished] == [ + (done[0].item.id, done[0].item.encrypted_content) + ] + message_indexes: Final = { + e.output_index for e in events if e.type == "response.output_item.added" and e.item.type == "message" + } + assert message_indexes and added[0].output_index not in message_indexes + assert completed[0].response.output_text == _answer(marker) + call, body = _only_call(wire, marker) + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + + +def _history(marker: str, compaction: Mapping[str, JsonValue]) -> list[JsonValue]: + return [ + {"role": "developer", "content": f"rules-{marker}"}, + {"role": "user", "content": f"old-{marker}"}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": f"old-answer-{marker}", "annotations": []}], + }, + dict(compaction), + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": f"recap-{marker}", "annotations": []}], + }, + {"role": "user", "content": _ask(marker)}, + ] + + +_ROUTES: Final[tuple[Route, ...]] = ( + "claude-compacting", + "claude-plain", + "gemini", + "openai", + "openai-compact", + "openai-opaque", +) + + +def _route_provider(route: Route) -> Literal["claude", "gemini", "openai"]: + if route.startswith("claude"): + return "claude" + return "gemini" if route == "gemini" else "openai" + + +def _assert_reached( + route: Route, marker: str, sent: Mapping[str, JsonValue], call: Request, body: Mapping[str, JsonValue] +) -> None: + summary: Final = _summary(marker) + dumped: Final = json.dumps(body) + if route == "openai-opaque": + assert urlsplit(call.target).path.endswith("/responses"), call.target + assert sent in _objects(body["input"]), body["input"] + assert f"old-{marker}" in dumped + return + if route != "claude-compacting": + assert f"old-{marker}" not in dumped and f"old-answer-{marker}" not in dumped, dumped[:2000] + assert f"rules-{marker}" in dumped + assert f"recap-{marker}" in dumped + match route: + case "claude-compacting": + assistant: Final = tuple( + message for message in _objects(body["messages"]) if message["role"] == "assistant" + ) + blocks: Final = tuple( + chain.from_iterable( + _objects(message["content"]) for message in assistant if isinstance(message["content"], list) + ) + ) + assert _block(summary) in blocks, body["messages"] + assert summary not in _anthropic_texts(body, "user") + assert f"rules-{marker}" in json.dumps(body["system"]) + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + case "claude-plain": + assert "compaction" not in _anthropic_block_types(body), body["messages"] + assert summary in _anthropic_texts(body, "user"), body["messages"] + case "gemini": + assert summary in _gemini_texts(body, "user"), body["contents"] + case "openai" | "openai-compact": + suffix: Final = "/responses/compact" if route == "openai-compact" else "/responses" + assert urlsplit(call.target).path.endswith(suffix), call.target + assert "compaction" not in _openai_item_types(body), body["input"] + assert summary in _openai_user_texts(body), body["input"] + + +@pytest.mark.parametrize("route", _ROUTES) +def test_compaction_item_reaches_each_route_in_the_shape_it_reads(gateway: Gateway, route: Route) -> None: + marker: Final = _new_marker() + compaction: Final = ( + {"type": "compaction", "id": f"cmp_{marker}", "encrypted_content": f"gAAAAA-opaque-{marker}"} + if route == "openai-opaque" + else _compaction_item(_summary(marker)) + ) + compacting: Final[dict[str, JsonValue]] = ( + {"context_management": _COMPACTING} if route == "claude-compacting" else {} + ) + dropping: Final[dict[str, JsonValue]] = ( + {"context_management": _COMPACTING, "drop_params": True} if route == "gemini" else {} + ) + path: Final = "/v1/responses/compact" if route == "openai-compact" else "/v1/responses" + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, _route_provider(route)) + response: Final = gateway.request( + "POST", + path, + {"model": model, "input": _history(marker, compaction), **compacting, **dropping, **_CACHE_BUST}, + ) + assert response.status_code == 200, response.text + if route != "openai-compact": + assert _answer(marker) in response.text, response.text + call, body = _only_call(wire, marker) + _assert_reached(route, marker, compaction, call, body) + + +def _input_tokens(gateway: Gateway, model: str, items: list[JsonValue]) -> int: + counted: Final = gateway.post("/v1/responses/input_tokens", {"model": model, "input": items}) + tokens: Final = counted["input_tokens"] + assert isinstance(tokens, int), counted + return tokens + + +def test_input_tokens_count_from_the_newest_compaction(gateway: Gateway) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + old_history: Final = " ".join(f"history-{index}" for index in range(600)) + compacted: Final[list[JsonValue]] = [ + {"role": "user", "content": old_history}, + _compaction_item(summary), + {"role": "user", "content": _ask(marker)}, + ] + equivalent: Final[list[JsonValue]] = [ + {"role": "user", "content": summary}, + {"role": "user", "content": _ask(marker)}, + ] + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + assert _input_tokens(gateway, model, compacted) == _input_tokens(gateway, model, equivalent) + counted: Final = tuple( + request.body for request in wire.drain() if urlsplit(request.target).path.endswith("/count_tokens") + ) + assert len(counted) == 2, counted + assert all(b"history-" not in body for body in counted), counted[0][:2000] + + +def test_streamed_compaction_item_is_done_before_the_reasoning_after_it(gateway: Gateway) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + events: Final = _stream( + gateway, + "/v1/responses", + { + "model": model, + "input": f"{_ask(marker, compact=True)} {_think(marker)}", + "context_management": _COMPACTING, + }, + ) + done: Final = tuple( + (position, event["output_index"], object_value(event["item"])["type"]) + for position, event in enumerate(events) + if event["type"] == "response.output_item.done" + ) + completed: Final = _only([event for event in events if event["type"] == "response.completed"]) + assert [(index, kind) for _, index, kind in done] == [(0, "compaction"), (1, "reasoning"), (2, "message")], [ + event["type"] for event in events + ] + assert [item["type"] for item in _objects(object_value(completed["response"])["output"])] == [ + "compaction", + "reasoning", + "message", + ] + reasoning_done: Final = done[1][0] + sequence_numbers: Final = [ + number for event in events[: reasoning_done + 1] if isinstance(number := event.get("sequence_number"), int) + ] + assert sequence_numbers == sorted(set(sequence_numbers)), sequence_numbers + _only_call(wire, marker) + + +def test_chat_stream_reports_the_whole_compaction_block(gateway: Gateway) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + chunks: Final = _stream( + gateway, + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _ask(marker, compact=True)}], + "context_management": {"edits": [{"type": "compact_20260112"}]}, + }, + ) + reported: Final = _streamed_compaction_blocks(chunks) + assert reported, chunks + assert _objects(reported[-1])[-1] == _block(_summary(marker)), reported[-1] + assert len(_provider_calls(wire, marker)) == 1 + + +def _guardrail_peer(request: Request) -> Reply: + payload: Final = _JSON_OBJECT.validate_json(request.body) + texts: Final = payload.get("texts") or [] + assert isinstance(texts, list), payload + if urlsplit(request.target).path.startswith("/block/"): + found: Final = any(isinstance(text, str) and _CITY in text for text in texts) + verdict: Final = {"action": "BLOCKED", "blocked_reason": "city in summary"} if found else {"action": "NONE"} + return Reply(body=json.dumps(verdict).encode()) + masked: Final = [text.replace(_CITY, _MASKED_CITY) if isinstance(text, str) else text for text in texts] + return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "texts": masked}).encode()) + + +def _guardrail(name: str, mode: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": api_base, + "api_key": "synthetic-guardrail-key", + **fields, + }, + } + + +def _deployment(name: str, model: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return {"model_name": name, "litellm_params": {"model": model, "api_base": api_base, **fields}} + + +_STORED_CLAUDE: Final = "compaction-stored-claude" +_SMALL_CLAUDE: Final = "compaction-small-claude" +_STORED_GEMINI: Final = "compaction-stored-gemini" +_STORED_OPENAI: Final = "compaction-stored-openai" +_SMALL_CONTEXT: Final = 2000 + + +def _stored_config(directory: Path, provider_url: str, guardrail_url: str, **litellm_settings: JsonValue) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + claude: Final = {"api_key": cc.ANTHROPIC_API_KEY} + config: Final[dict[str, JsonValue]] = { + **stock, + "model_list": [ + _deployment(_STORED_CLAUDE, f"anthropic/{_CLAUDE}", provider_url, **claude), + { + **_deployment(_SMALL_CLAUDE, f"anthropic/{_CLAUDE}", provider_url, **claude), + "model_info": {"max_input_tokens": _SMALL_CONTEXT}, + }, + _deployment(_STORED_GEMINI, f"gemini/{_GEMINI}", provider_url, api_key="synthetic-gemini-key"), + _deployment(_STORED_OPENAI, f"openai/{_OPENAI}", f"{provider_url}/v1", api_key="synthetic-openai-key"), + ], + "general_settings": { + **object_value(stock["general_settings"]), + "store_prompts_in_spend_logs": True, + "disable_responses_id_security": True, + }, + "litellm_settings": {**object_value(stock["litellm_settings"]), **litellm_settings}, + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0, "enable_pre_call_checks": True}, + "guardrails": [ + _guardrail(_MASK_INPUT, "pre_call", f"{guardrail_url}/input"), + _guardrail(_MASK_OUTPUT, "post_call", f"{guardrail_url}/output"), + _guardrail(_BLOCK_OUTPUT, "post_call", f"{guardrail_url}/block", streaming_buffer_until_moderated=True), + ], + } + path: Final = directory / "compaction-stored.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@dataclass(frozen=True, slots=True) +class _StoredRig: + gateway: Gateway + provider: Wire + guardrail: Wire + + +def _stored_rig(directory: Path, **litellm_settings: JsonValue) -> Iterator[_StoredRig]: + with ( + wire_server(_vendor) as provider, + wire_server(_guardrail_peer) as guardrail, + gateway_from_environment() as gateway, + ): + config: Final = _stored_config(directory, provider.url, guardrail.url, **litellm_settings) + with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: + yield _StoredRig(owned, provider, guardrail) + + +@pytest.fixture(scope="module") +def stored(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + yield from _stored_rig(tmp_path_factory.mktemp("compaction-stored")) + + +@pytest.fixture(scope="module") +def redacting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + yield from _stored_rig(tmp_path_factory.mktemp("compaction-redacting"), turn_off_message_logging=True) + + +def _logged_row(request_id: str, status: Literal["success", "failure"]) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, proxy_server_request, response, metadata, status FROM "LiteLLM_SpendLogs"' + " WHERE request_id=%s", + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["status"] == status, rows[0] + return rows[0] + + +def _spend_row(response_id: str) -> dict[str, JsonValue]: + decoded: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id).get("response_id") + return _logged_row(decoded if isinstance(decoded, str) else response_id, "success") + + +_REPLAY_ROUTES: Final = ( + pytest.param(_STORED_CLAUDE, True, id="claude-compacting"), + pytest.param(_STORED_CLAUDE, False, id="claude-plain"), + pytest.param(_STORED_GEMINI, True, id="gemini"), +) + + +@pytest.mark.parametrize(("follow_up_model", "compacting"), _REPLAY_ROUTES) +def test_previous_response_id_after_a_compaction_turn_replays_the_whole_summary( + stored: _StoredRig, follow_up_model: str, compacting: bool +) -> None: + first: Final = _new_marker() + second: Final = _new_marker() + summary: Final = _summary(first) + long_text: Final = f"long-{first} {'z' * 3000} tail-{first}" + compaction_request: Final[dict[str, JsonValue]] = {"context_management": _COMPACTING} + opened: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": [{"role": "user", "content": long_text}, {"role": "user", "content": _ask(first, compact=True)}], + **compaction_request, + }, + ) + assert _decoded(string_value(_only(_output_items(opened, "compaction"))["encrypted_content"])) == _block(summary) + row: Final = _spend_row(string_value(opened["id"])) + assert summary in json.dumps(row["response"]), json.dumps(row["response"])[:600] + stored_request: Final = json.dumps(row["proxy_server_request"]) + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored_request and "z" * 3000 not in stored_request + _only_call(stored.provider, first) + + follow_up: Final = _responses( + stored.gateway, + { + "model": follow_up_model, + "input": _ask(second), + "previous_response_id": opened["id"], + **(compaction_request if compacting else {}), + **({"drop_params": True} if follow_up_model == _STORED_GEMINI else {}), + }, + ) + assert _answer(second) in json.dumps(follow_up) + call, body = _only_call(stored.provider, second) + dumped: Final = json.dumps(body) + assert f"long-{first}" not in dumped and f"ask-{first}" not in dumped, dumped[:2000] + assert _answer(first) in dumped + if follow_up_model == _STORED_GEMINI: + assert summary in _gemini_texts(body, "user"), body["contents"] + elif compacting: + replayed: Final = _objects(body["messages"])[0] + assert replayed["role"] == "assistant", body["messages"] + assert _objects(replayed["content"])[0] == _block(summary), replayed + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + else: + assert "compaction" not in _anthropic_block_types(body), body["messages"] + assert summary in _anthropic_texts(body, "user"), body["messages"] + + +def test_router_pre_call_check_counts_from_the_newest_compaction(stored: _StoredRig) -> None: + compacted_marker: Final = _new_marker() + control_marker: Final = _new_marker() + old_history: Final = " ".join(f"history-{index}" for index in range(3 * _SMALL_CONTEXT)) + + def send(marker: str, items: list[JsonValue]) -> httpx.Response: + return stored.gateway.request( + "POST", + "/v1/responses", + {"model": _SMALL_CLAUDE, "input": [*items, {"role": "user", "content": _ask(marker)}], **_CACHE_BUST}, + ) + + compacted: Final = send( + compacted_marker, + [{"role": "user", "content": old_history}, _compaction_item(f"short summary {compacted_marker}")], + ) + assert compacted.status_code == 200, compacted.text + _only_call(stored.provider, compacted_marker) + control: Final = send(control_marker, [{"role": "user", "content": old_history}]) + assert control.status_code == 400, control.text + assert _provider_calls(stored.provider, control_marker) == () + + +_GUARDED_ROUTES: Final = ( + pytest.param(_STORED_OPENAI, id="openai"), + pytest.param(_STORED_CLAUDE, id="claude-compacting"), +) + + +@pytest.mark.parametrize("model", _GUARDED_ROUTES) +def test_pre_call_guardrail_rewrites_the_compaction_summary(stored: _StoredRig, model: str) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + masked: Final = summary.replace(_CITY, _MASKED_CITY) + compacting: Final[dict[str, JsonValue]] = {"context_management": _COMPACTING} if model == _STORED_CLAUDE else {} + _responses( + stored.gateway, + { + "model": model, + "input": [_compaction_item(summary), {"role": "user", "content": _ask(marker)}], + "guardrails": [_MASK_INPUT], + **compacting, + }, + ) + scanned: Final = tuple( + _JSON_OBJECT.validate_json(request.body) + for request in stored.guardrail.drain() + if f"kept-{marker}".encode() in request.body + ) + assert any(summary in _listed(payload["texts"]) for payload in scanned), scanned + call, body = _only_call(stored.provider, marker) + assert _CITY not in json.dumps(body), json.dumps(body)[:2000] + if model == _STORED_OPENAI: + assert masked in _openai_user_texts(body), body["input"] + else: + assert _objects(_objects(body["messages"])[0]["content"])[0] == _block(masked), body["messages"] + + +def test_post_call_guardrail_rewrites_the_returned_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + returned: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": _ask(marker, compact=True), + "context_management": _COMPACTING, + "guardrails": [_MASK_OUTPUT], + }, + ) + masked: Final = _block(_summary(marker).replace(_CITY, _MASKED_CITY)) + assert _decoded(string_value(_only(_output_items(returned, "compaction"))["encrypted_content"])) == masked + _only_call(stored.provider, marker) + + +def test_buffered_stream_guardrail_blocks_on_the_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + body: Final[dict[str, JsonValue]] = { + "model": _STORED_CLAUDE, + "input": _ask(marker, compact=True), + "context_management": _COMPACTING, + "guardrails": [_BLOCK_OUTPUT], + "stream": True, + **_CACHE_BUST, + } + headers: Final = {"Authorization": f"Bearer {stored.gateway.key}"} + with stored.gateway.client.stream("POST", "/v1/responses", json=body, headers=headers) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + peer_calls: Final = tuple( + (urlsplit(request.target).path, _JSON_OBJECT.validate_json(request.body).get("texts")) + for request in stored.guardrail.drain() + if marker.encode() in request.body + ) + assert any(path.startswith("/block/") and _summary(marker) in _listed(texts) for path, texts in peer_calls), [ + (path, [str(value)[:80] for value in _listed(texts or [])]) for path, texts in peer_calls + ] + event_types: Final = tuple(event["type"] for event in _sse_events(text)) + assert "response.failed" in event_types, event_types + assert "response.completed" not in event_types, event_types + assert f"kept-{marker}" not in text and _CITY not in text, text[:2000] + _only_call(stored.provider, marker) + + +def test_redacted_spend_row_holds_no_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + returned: Final = _responses( + stored.gateway, + {"model": _STORED_CLAUDE, "input": _ask(marker, compact=True), "context_management": _COMPACTING}, + headers=_REDACT_HEADER, + ) + assert _decoded(string_value(_only(_output_items(returned, "compaction"))["encrypted_content"])) == _block( + _summary(marker) + ) + row: Final = _spend_row(string_value(returned["id"])) + assert f"kept-{marker}" not in json.dumps(row["response"]), json.dumps(row["response"])[:2000] + _only_call(stored.provider, marker) + + +def test_stored_compaction_input_item_keeps_its_whole_summary_and_nothing_else(stored: _StoredRig) -> None: + first: Final = _new_marker() + second: Final = _new_marker() + summary: Final = _summary(first) + padded_item: Final = { + **_compaction_item(summary), + "encrypted_content": json.dumps({**_block(summary), "padding": f"pad-{first} {'p' * 3000}"}), + } + opened: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": [padded_item, {"role": "user", "content": _ask(first)}], + "context_management": _COMPACTING, + }, + ) + stored_input: Final = _objects( + object_value(_spend_row(string_value(opened["id"]))["proxy_server_request"])["input"] + ) + assert _decoded(string_value(stored_input[0]["encrypted_content"])) == _block(summary), stored_input[0] + _only_call(stored.provider, first) + + _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": _ask(second), + "previous_response_id": opened["id"], + "context_management": _COMPACTING, + }, + ) + _, body = _only_call(stored.provider, second) + assert summary in _anthropic_texts(body, "user"), body["messages"] + assert f"pad-{first}" not in json.dumps(body), body["messages"] + + +def test_redacted_chat_stream_spend_row_holds_no_compaction_delta(stored: _StoredRig) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + chunks: Final = _stream( + stored.gateway, + "/v1/chat/completions", + { + "model": _STORED_CLAUDE, + "messages": [{"role": "user", "content": _ask(marker, compact=True)}], + "context_management": {"edits": [{"type": "compact_20260112"}]}, + }, + headers=_REDACT_HEADER, + ) + streamed_blocks: Final = _streamed_compaction_blocks(chunks) + assert streamed_blocks and _objects(streamed_blocks[-1])[-1] == _block(summary), chunks + secrets: Final = (f"kept-{marker}", summary[len(summary) // 2 :]) + logged: Final = {column: json.dumps(value) for column, value in _spend_row(string_value(chunks[0]["id"])).items()} + leaked: Final = {column: text for column, text in logged.items() if any(secret in text for secret in secrets)} + assert not leaked, {column: text[max(0, text.find(secrets[0]) - 600) :][:1500] for column, text in leaked.items()} + _only_call(stored.provider, marker) + + +def _failing_request(path: str, marker: str) -> dict[str, JsonValue]: + ask: Final = f"{_ask(marker)} {_fail(marker)}" + if path == "/v1/responses": + return { + "model": _STORED_CLAUDE, + "input": [_compaction_item(_summary(marker)), {"role": "user", "content": ask}], + } + return {"model": _STORED_CLAUDE, "max_tokens": 64, "messages": [{"role": "user", "content": ask}]} + + +def _failed_call_row( + rig: _StoredRig, + path: str, + marker: str, + headers: Mapping[str, str], + key: str | None = None, + **fields: JsonValue, +) -> dict[str, JsonValue]: + body: Final = {**_failing_request(path, marker), **fields} + response: Final = rig.gateway.request("POST", path, body, headers=headers, key=key) + assert response.status_code == 400, response.text + _only_call(rig.provider, marker) + return _logged_row(response.headers["x-litellm-call-id"], "failure") + + +_LEGACY_REDACT_HEADER: Final = {"litellm-enable-message-redaction": "true"} +_REDACTED_FAILURES: Final = ( + pytest.param("/v1/responses", _REDACT_HEADER, id="responses"), + pytest.param("/v1/responses", _LEGACY_REDACT_HEADER, id="responses-legacy-header"), + pytest.param("/v1/messages", _REDACT_HEADER, id="messages"), + pytest.param("/v1/chat/completions", _REDACT_HEADER, id="chat"), +) + + +@pytest.mark.parametrize(("path", "headers"), _REDACTED_FAILURES) +def test_failed_call_spend_row_holds_no_request_text_with_the_redaction_header( + stored: _StoredRig, path: str, headers: Mapping[str, str] +) -> None: + marker: Final = _new_marker() + logged: Final = json.dumps(_failed_call_row(stored, path, marker, headers)) + assert marker not in logged, logged[max(0, logged.find(marker) - 600) :][:1500] + + +def test_failed_responses_call_spend_row_keeps_the_whole_compaction_summary_without_the_header( + stored: _StoredRig, +) -> None: + marker: Final = _new_marker() + row: Final = _failed_call_row(stored, "/v1/responses", marker, {}) + stored_input: Final = _objects(object_value(row["proxy_server_request"])["input"]) + assert _decoded(string_value(stored_input[0]["encrypted_content"])) == _block(_summary(marker)), stored_input[0] + assert f"{_ask(marker)} {_fail(marker)}" in _message_texts(stored_input[1]), stored_input[1] + + +def test_failed_responses_call_stores_the_caller_metadata_without_the_request_headers(stored: _StoredRig) -> None: + marker: Final = _new_marker() + row: Final = _failed_call_row(stored, "/v1/responses", marker, _REDACT_HEADER, metadata={"caller": marker}) + stored_metadata: Final = object_value(object_value(row["proxy_server_request"])["metadata"]) + assert stored_metadata.get("caller") == marker, stored_metadata + assert "headers" not in stored_metadata, stored_metadata + + +def test_failed_responses_call_keeps_the_caller_metadata_headers_over_the_redaction_header(stored: _StoredRig) -> None: + marker: Final = _new_marker() + caller_headers: Final = {"caller": marker} + row: Final = _failed_call_row(stored, "/v1/responses", marker, _REDACT_HEADER, metadata={"headers": caller_headers}) + stored_request: Final = object_value(row["proxy_server_request"]) + assert object_value(stored_request["metadata"])["headers"] == caller_headers, stored_request["metadata"] + assert _ask(marker) in json.dumps(stored_request), json.dumps(stored_request)[:1500] + + +_GLOBALLY_REDACTED_FAILURES: Final = ( + pytest.param("/v1/responses", _UNREDACT_HEADER, False, id="responses-disable-header"), + pytest.param("/v1/chat/completions", _UNREDACT_HEADER, True, id="chat-disable-header"), + pytest.param("/v1/responses", {}, False, id="responses-no-header"), +) + + +@pytest.mark.parametrize(("path", "headers", "kept"), _GLOBALLY_REDACTED_FAILURES) +def test_opted_out_key_failed_call_spend_row_takes_the_disable_header_over_global_redaction_only_on_chat( + redacting: _StoredRig, path: str, headers: Mapping[str, str], kept: bool +) -> None: + marker: Final = _new_marker() + with redacting.gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"allow_client_message_redaction_opt_out": True}) + row: Final = _failed_call_row(redacting, path, marker, headers, key) + stored_request: Final = json.dumps(row["proxy_server_request"]) + assert (_ask(marker) in stored_request) is kept, stored_request[:1500] diff --git a/tests/integration/sdk/test_dual_cache_redis.py b/tests/integration/sdk/test_dual_cache_redis.py index ddd01480d36..3527dbf91ee 100644 --- a/tests/integration/sdk/test_dual_cache_redis.py +++ b/tests/integration/sdk/test_dual_cache_redis.py @@ -6,6 +6,7 @@ from unittest.mock import patch import pytest +from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache @@ -92,3 +93,66 @@ async def test_sync_and_async_batch_reads_share_one_redis_without_sync_reads_goi assert in_loop_results == [[expected[key]] for key in in_loop_keys] assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]] + + +@pytest.mark.asyncio +async def test_concurrent_session_repin_claims_keep_the_shared_winner_after_recovery() -> None: + redis: Final = _redis_cache() + workers: Final = (DualCache(redis_cache=redis), DualCache(redis_cache=redis)) + key: Final = f"session-repin-{uuid.uuid4().hex}" + choices: Final = ({"model_id": "replacement"}, {"model_id": "other"}) + try: + await redis.async_set_cache(key, {"model_id": "original"}, ttl=60) + for worker in workers: + worker.in_memory_cache.set_cache(key, {"model_id": "original"}, ttl=60) + winners: Final = await asyncio.gather( + *( + claim_affinity_pin(worker, key, choice, 60, eligible_values=choices, preserve_legacy_model_ids=True) + for worker, choice in zip(workers, choices, strict=True) + ) + ) + assert winners[0] == winners[1] and winners[0] in choices + late_original: Final = await claim_affinity_pin( + workers[0], + key, + {"model_id": "original"}, + 60, + eligible_values=({"model_id": "original"}, *choices), + preserve_legacy_model_ids=True, + ) + assert late_original == winners[0] + assert await redis.async_get_cache(key) == winners[0] + finally: + await redis.async_delete_cache(key) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored", ["original", '"original"', '{"model_id":"original"}']) +async def test_session_repin_lua_preserves_healthy_legacy_pins_and_replaces_unavailable_ones(stored: str) -> None: + redis: Final = _redis_cache() + cache: Final = DualCache(redis_cache=redis) + key: Final = f"session-repin-legacy-{uuid.uuid4().hex}" + replacement: Final = {"model_id": "replacement"} + try: + redis.redis_client.set(redis.check_and_fix_namespace(key), stored, ex=60) + preserved: Final = await claim_affinity_pin( + cache, + key, + replacement, + 60, + eligible_values=({"model_id": "original"}, replacement), + preserve_legacy_model_ids=True, + ) + assert preserved in ("original", {"model_id": "original"}) + moved: Final = await claim_affinity_pin( + cache, + key, + replacement, + 60, + eligible_values=(replacement,), + preserve_legacy_model_ids=True, + ) + assert moved == replacement + assert await redis.async_get_cache(key) == replacement + finally: + await redis.async_delete_cache(key) diff --git a/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py b/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py new file mode 100644 index 00000000000..dc94c96da05 --- /dev/null +++ b/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py @@ -0,0 +1,474 @@ +from __future__ import annotations + +import asyncio +import json +import linecache +import sys +import threading +import uuid +from collections.abc import AsyncIterable, Awaitable, Callable, Iterable +from concurrent.futures import ThreadPoolExecutor +from types import FrameType +from typing import Final, TypeVar + +import pytest +from integration._support.client import object_value +from integration._support.openai_wire import chat_reply, openai_error, responses_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +import litellm +from litellm import Router + +_DEPLOYMENTS: Final = 4 +_TARGET: Final = _DEPLOYMENTS - 1 +_SYNC_SECONDS: Final = 10 +_RATE_LIMITED_KEY: Final = "key-rate-limited" +_BURST_THREADS: Final = 8 +_BURST_CALLS: Final = 24 +_CHURN_UPSERTS: Final = 120 +_UNREPORTED_MODEL_ID: Final = frozenset({"aresponses-stream", "aanthropic-messages"}) +_ROUTER_FILE: Final = Router.__init__.__code__.co_filename +_INDEX_LOCK_LINES: Final = frozenset( + number + for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1) + if "with self._deployment_index_lock" in line +) + +T = TypeVar("T") + + +def _peer(request: Request) -> Reply: + if request.headers["authorization"] == f"Bearer {_RATE_LIMITED_KEY}": + return openai_error(429) + body: Final = object_value(json.loads(request.body)) + stream: Final = body.get("stream") is True + if request.target == "/v1/responses": + return responses_reply(f"resp_{uuid.uuid4().hex}", "gpt-4o-mini", "ok", stream=stream) + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "gpt-4o-mini", "ok", stream=stream) + + +def _deployment( + index: int, + api_base: str, + team: bool, + *, + model: int | None = None, + public_name: str | None = None, + timeout: float = 10, +) -> dict[str, JsonValue]: + team_info: Final[dict[str, JsonValue]] = ( + {"team_id": "team-1", "team_public_model_name": public_name or f"public-{index}"} if team else {} + ) + return { + "model_name": f"model-{index if model is None else model}", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": api_base, + "api_key": f"key-{index}", + "timeout": timeout, + }, + "model_info": {"id": f"deployment-{index}", **team_info}, + } + + +def _router(wire: Wire, team: bool = False) -> Router: + return Router(model_list=[_deployment(index, f"{wire.url}/v1", team) for index in range(_DEPLOYMENTS)]) + + +def _read_during_held_write(read: Callable[[], T], write: Callable[[], object]) -> T: + """Run `read` on its own thread while `write` waits where the list has shrunk but the index maps have not caught + up. The write resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets + the half-updated list and a read that takes the lock waits for the write""" + reindex: Final = Router._update_deployment_indices_after_removal.__code__ + held: Final = threading.Event() + release: Final = threading.Event() + settled: Final = threading.Event() + released_in_time: Final[list[bool]] = [] + outcome: Final[list[T]] = [] + failure: Final[list[BaseException]] = [] + + def hold_at_reindex(frame: FrameType, event: str, _arg: object) -> None: + if event == "call" and frame.f_code is reindex and not held.is_set(): + held.set() + released_in_time.append(release.wait(timeout=_SYNC_SECONDS)) + + def run_write() -> None: + previous: Final = sys.getprofile() + sys.setprofile(hold_at_reindex) + try: + write() + finally: + sys.setprofile(previous) + + def watch_read(frame: FrameType, event: str, _arg: object) -> Callable[..., object] | None: + if frame.f_code.co_filename != _ROUTER_FILE: + return None + if event == "line" and frame.f_lineno in _INDEX_LOCK_LINES: + settled.set() + return watch_read + + def run_read() -> None: + previous: Final = sys.gettrace() + sys.settrace(watch_read) + try: + outcome.append(read()) + except BaseException as error: + failure.append(error) + finally: + sys.settrace(previous) + settled.set() + + writer: Final = threading.Thread(target=run_write) + reader: Final = threading.Thread(target=run_read) + writer.start() + try: + assert held.wait(timeout=_SYNC_SECONDS), "the write never reached re-indexing" + reader.start() + assert settled.wait(timeout=_SYNC_SECONDS), "the read neither returned nor reached the index lock" + finally: + release.set() + if reader.ident is not None: + reader.join() + writer.join() + assert released_in_time == [True], "the write was not held for the whole read" + if failure: + raise failure[0] + return outcome[0] + + +def _read_during_held_delete(router: Router, deployment_id: str, read: Callable[[], T]) -> T: + return _read_during_held_write(read, lambda: router.delete_deployment(deployment_id)) + + +def _keys(wire: Wire) -> tuple[str, ...]: + return tuple(request.headers["authorization"] for request in wire.drain()) + + +def _team_metadata(team: bool) -> dict[str, JsonValue]: + return {"user_api_key_team_id": "team-1"} if team else {} + + +def _completion(router: Router, model: str, team: bool) -> str | None: + response: Final = router.completion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team) + ) + return response._hidden_params["model_id"] + + +def _completion_stream(router: Router, model: str, team: bool) -> str | None: + stream: Final = router.completion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team), stream=True + ) + assert isinstance(stream, Iterable), stream + chunks: Final = list(stream) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok", chunks + return stream._hidden_params["model_id"] + + +def _acompletion(router: Router, model: str, team: bool) -> str | None: + async def call() -> str | None: + response: Final = await router.acompletion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team) + ) + return response._hidden_params["model_id"] + + return asyncio.run(call()) + + +def _acompletion_stream(router: Router, model: str, team: bool) -> str | None: + async def call() -> str | None: + stream: Final = await router.acompletion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team), stream=True + ) + assert isinstance(stream, AsyncIterable), stream + chunks: Final = [chunk async for chunk in stream] + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok", chunks + return stream._hidden_params["model_id"] + + return asyncio.run(call()) + + +def _responses(router: Router, model: str, team: bool) -> str | None: + response: Final = router.responses(model=model, input="hi", metadata=_team_metadata(team)) + assert response.output_text == "ok", response + return response._hidden_params["model_id"] + + +def _aresponses(router: Router, model: str, team: bool) -> str | None: + async def call() -> str | None: + response: Final = await router.aresponses(model=model, input="hi", metadata=_team_metadata(team)) + assert response.output_text == "ok", response + return response._hidden_params["model_id"] + + return asyncio.run(call()) + + +def _aresponses_stream(router: Router, model: str, team: bool) -> str | None: + async def call() -> None: + stream: Final = await router.aresponses(model=model, input="hi", metadata=_team_metadata(team), stream=True) + assert isinstance(stream, AsyncIterable), stream + events: Final = [event.type async for event in stream] + assert events[-1] == "response.completed", events + + asyncio.run(call()) + return None + + +def _aanthropic_messages(router: Router, model: str, team: bool) -> str | None: + async def call() -> None: + response: Final = await router.aanthropic_messages( + model=model, messages=[{"role": "user", "content": "hi"}], max_tokens=16, metadata=_team_metadata(team) + ) + assert response["content"][0]["text"] == "ok", response + + asyncio.run(call()) + return None + + +_CALLS: Final[dict[str, Callable[[Router, str, bool], str | None]]] = { + "completion": _completion, + "completion-stream": _completion_stream, + "acompletion": _acompletion, + "acompletion-stream": _acompletion_stream, + "responses": _responses, + "aresponses": _aresponses, + "aresponses-stream": _aresponses_stream, + "aanthropic-messages": _aanthropic_messages, +} + + +@pytest.mark.parametrize("call", list(_CALLS)) +@pytest.mark.parametrize("team", [False, True], ids=["model-name", "team-public-name"]) +def test_call_reaches_its_deployment_while_another_thread_deletes_one(call: str, team: bool) -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire, team) + served_by: Final = _read_during_held_delete( + router, + "deployment-0", + lambda: _CALLS[call](router, f"public-{_TARGET}" if team else f"model-{_TARGET}", team), + ) + keys: Final = _keys(wire) + + assert served_by == (None if call in _UNREPORTED_MODEL_ID else f"deployment-{_TARGET}") + assert keys == (f"Bearer key-{_TARGET}",) + assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)] + + +@pytest.mark.parametrize("deleted", [0, _TARGET - 1], ids=["first", "next-to-target"]) +def test_completion_reaches_its_deployment_whichever_earlier_deployment_is_deleted(deleted: int) -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + served_by: Final = _read_during_held_delete( + router, f"deployment-{deleted}", lambda: _completion(router, f"model-{_TARGET}", team=False) + ) + keys: Final = _keys(wire) + + assert served_by == f"deployment-{_TARGET}" + assert keys == (f"Bearer key-{_TARGET}",) + + +def test_team_completion_reaches_the_remaining_deployment_of_a_shared_public_name_while_its_sibling_is_deleted() -> ( + None +): + with wire_server(_peer) as wire: + api_base: Final = f"{wire.url}/v1" + router: Final = Router( + model_list=[ + _deployment(0, api_base, team=False), + _deployment(1, api_base, team=False), + _deployment(2, api_base, team=True, public_name="public-shared"), + _deployment(3, api_base, team=True, public_name="public-shared"), + ] + ) + served_by: Final = _read_during_held_delete( + router, "deployment-2", lambda: _completion(router, "public-shared", team=True) + ) + keys: Final = _keys(wire) + + assert served_by == "deployment-3" + assert keys == ("Bearer key-3",) + + +def test_completion_for_the_deployment_being_deleted_fails_cleanly_and_sends_nothing() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + with pytest.raises(litellm.BadRequestError): + _read_during_held_delete( + router, + "deployment-0", + lambda: router.completion(model="model-0", messages=[{"role": "user", "content": "hi"}]), + ) + keys: Final = _keys(wire) + + assert keys == () + assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)] + + +def test_provider_error_during_a_concurrent_delete_comes_from_the_requested_deployment() -> None: + with wire_server(_peer) as wire: + api_base: Final = f"{wire.url}/v1" + rate_limited: Final = _deployment(_TARGET, api_base, team=False) + rate_limited["litellm_params"] = {**object_value(rate_limited["litellm_params"]), "api_key": _RATE_LIMITED_KEY} + router: Final = Router( + model_list=[*(_deployment(index, api_base, team=False) for index in range(_TARGET)), rate_limited], + num_retries=0, + ) + with pytest.raises(litellm.RateLimitError): + _read_during_held_delete( + router, + "deployment-0", + lambda: router.completion(model=f"model-{_TARGET}", messages=[{"role": "user", "content": "hi"}]), + ) + served_by: Final = _completion(router, "model-1", team=False) + keys: Final = _keys(wire) + + assert served_by == "deployment-1" + assert keys == (f"Bearer {_RATE_LIMITED_KEY}", "Bearer key-1") + + +def test_two_threads_deleting_the_same_deployment_remove_it_once() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + second: Final = _read_during_held_delete( + router, "deployment-0", lambda: router.delete_deployment("deployment-0") + ) + served_by: Final = _completion(router, f"model-{_TARGET}", team=False) + keys: Final = _keys(wire) + + assert second is None + assert served_by == f"deployment-{_TARGET}" + assert keys == (f"Bearer key-{_TARGET}",) + assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)] + + +def test_completion_reaches_its_deployment_while_another_thread_upserts_an_earlier_one() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + changed: Final = litellm.Deployment(**_deployment(0, f"{wire.url}/v1", team=False, timeout=11)) + served_by: Final = _read_during_held_write( + lambda: _completion(router, f"model-{_TARGET}", team=False), lambda: router.upsert_deployment(changed) + ) + keys: Final = _keys(wire) + + assert served_by == f"deployment-{_TARGET}" + assert keys == (f"Bearer key-{_TARGET}",) + assert sorted(router.get_model_ids()) == [f"deployment-{index}" for index in range(_DEPLOYMENTS)] + + +def test_deployment_added_during_a_concurrent_delete_serves_its_own_requests() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + added: Final = litellm.Deployment(**_deployment(9, f"{wire.url}/v1", team=False)) + _read_during_held_delete(router, "deployment-0", lambda: router.add_deployment(added)) + served: Final = tuple(_completion(router, f"model-{index}", team=False) for index in (9, _TARGET)) + keys: Final = _keys(wire) + + assert served == ("deployment-9", f"deployment-{_TARGET}") + assert keys == ("Bearer key-9", f"Bearer key-{_TARGET}") + + +def test_model_list_set_during_a_concurrent_delete_serves_its_own_requests() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + replacement: Final = [_deployment(index, f"{wire.url}/v1", team=False) for index in range(10, 14)] + _read_during_held_delete(router, "deployment-0", lambda: router.set_model_list(replacement)) + served: Final = tuple(_completion(router, f"model-{index}", team=False) for index in range(10, 14)) + keys: Final = _keys(wire) + + assert served == tuple(f"deployment-{index}" for index in range(10, 14)) + assert keys == tuple(f"Bearer key-{index}" for index in range(10, 14)) + assert router.get_model_ids() == [f"deployment-{index}" for index in range(10, 14)] + + +def _burst_router(api_base: str) -> Router: + """Two deployments per model name, so an upsert of one leaves its model group served.""" + return Router( + model_list=[ + _deployment(index, api_base, team=False, model=index % _DEPLOYMENTS) for index in range(2 * _DEPLOYMENTS) + ] + ) + + +def _churn(router: Router, api_base: str, stop: threading.Event) -> int: + """Upsert the deployments in turn with a changed timeout; each upsert moves one to the end of the list.""" + upserts = 0 + while upserts < _CHURN_UPSERTS and not stop.is_set(): + index, round_number = upserts % (2 * _DEPLOYMENTS), upserts // (2 * _DEPLOYMENTS) + router.upsert_deployment( + litellm.Deployment( + **_deployment(index, api_base, team=False, model=index % _DEPLOYMENTS, timeout=11 + round_number % 2) + ) + ) + upserts += 1 + return upserts + + +def _served_keys(wire: Wire) -> dict[str, str]: + """Which key reached the upstream for each request marker.""" + return { + object_value(json.loads(request.body))["messages"][0]["content"]: request.headers["authorization"] + for request in wire.drain() + } + + +def _assert_each_request_used_its_models_deployment(served: list[tuple[str, str | None]], keys: dict[str, str]) -> None: + for marker, model_id in served: + model: Final = int(marker.split(":")[0]) + assert model_id in (f"deployment-{model}", f"deployment-{model + _DEPLOYMENTS}"), (marker, model_id) + assert keys[marker] == f"Bearer key-{str(model_id).removeprefix('deployment-')}", (marker, keys[marker]) + assert len(keys) == len(served) + + +def test_thread_burst_during_upsert_churn_sends_every_request_with_its_deployments_key() -> None: + with wire_server(_peer) as wire: + api_base: Final = f"{wire.url}/v1" + router: Final = _burst_router(api_base) + stop: Final = threading.Event() + + def call(number: int) -> tuple[str, str | None]: + marker: Final = f"{number % _DEPLOYMENTS}:{number}" + response: Final = router.completion( + model=f"model-{number % _DEPLOYMENTS}", messages=[{"role": "user", "content": marker}] + ) + return marker, response._hidden_params["model_id"] + + with ThreadPoolExecutor(max_workers=_BURST_THREADS + 1) as pool: + churn: Final = pool.submit(_churn, router, api_base, stop) + served: Final = list(pool.map(call, range(_BURST_THREADS * _BURST_CALLS))) + stop.set() + upserts: Final = churn.result() + keys: Final = _served_keys(wire) + + assert upserts > 0 + _assert_each_request_used_its_models_deployment(served, keys) + assert sorted(router.get_model_ids()) == sorted(f"deployment-{index}" for index in range(2 * _DEPLOYMENTS)) + + +async def _gathered(calls: Iterable[Awaitable[tuple[str, str | None]]]) -> list[tuple[str, str | None]]: + return list(await asyncio.gather(*calls)) + + +def test_async_burst_during_upsert_churn_sends_every_request_with_its_deployments_key() -> None: + with wire_server(_peer) as wire: + api_base: Final = f"{wire.url}/v1" + router: Final = _burst_router(api_base) + stop: Final = threading.Event() + + async def call(number: int) -> tuple[str, str | None]: + marker: Final = f"{number % _DEPLOYMENTS}:{number}" + response: Final = await router.acompletion( + model=f"model-{number % _DEPLOYMENTS}", messages=[{"role": "user", "content": marker}] + ) + return marker, response._hidden_params["model_id"] + + with ThreadPoolExecutor(max_workers=1) as pool: + churn: Final = pool.submit(_churn, router, api_base, stop) + served: Final = asyncio.run(_gathered(call(number) for number in range(_BURST_THREADS * _BURST_CALLS))) + stop.set() + upserts: Final = churn.result() + keys: Final = _served_keys(wire) + + assert upserts > 0 + _assert_each_request_used_its_models_deployment(served, keys) + assert sorted(router.get_model_ids()) == sorted(f"deployment-{index}" for index in range(2 * _DEPLOYMENTS)) diff --git a/tests/integration/sdk/test_sync_stream_connection_release_wire.py b/tests/integration/sdk/test_sync_stream_connection_release_wire.py index a8b49d285c3..21774df66d0 100644 --- a/tests/integration/sdk/test_sync_stream_connection_release_wire.py +++ b/tests/integration/sdk/test_sync_stream_connection_release_wire.py @@ -212,3 +212,36 @@ def test_sync_stream_that_raises_mid_stream_releases_the_connection() -> None: assert _open_sessions_settling_to(wire, at_most=0) == 0, ( "a stream that raised mid-way left its TCP session open" ) + + +_STUCK: Final = _Shape( + "hosted_vllm", + "hosted_vllm/stuck-model", + _OPENAI_OPENING, + (_openai_delta("the model is stuck on this", None),) * (10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT), +) + + +def _still_generating_peer(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=(_STUCK.opening, *_STUCK.closing), pause_between_chunks=0.01) + + +def test_sync_stream_stuck_repeating_one_chunk_releases_the_connection() -> None: + with wire_server(_still_generating_peer, keep_alive=True) as wire: + stuck: Final = litellm.completion(**_request_kwargs(_STUCK, wire)) + with pytest.raises(Exception, match="repeating the same chunk"): + for _ in stuck: + pass + assert _open_sessions_settling_to(wire, at_most=0) == 0, "a stream stuck on one chunk left its TCP session open" + + +@pytest.mark.asyncio +async def test_async_stream_stuck_repeating_one_chunk_releases_the_connection() -> None: + with wire_server(_still_generating_peer, keep_alive=True) as wire: + stuck: Final = await litellm.acompletion(**_request_kwargs(_STUCK, wire)) + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in stuck: + pass + assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=0) == 0, ( + "a stream stuck on one chunk left its TCP session open" + ) diff --git a/tests/integration/security/test_offline_mode.py b/tests/integration/security/test_offline_mode.py new file mode 100644 index 00000000000..11012eb14bc --- /dev/null +++ b/tests/integration/security/test_offline_mode.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import json +import uuid +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, object_value +from integration._support.forward_proxy import refusing_forward_proxy +from integration._support.process import owned_proxy +from pydantic import JsonValue + +_PACKAGE: Final = Path(__file__).resolve().parents[3] / "litellm" +_CONTROL_HOST: Final = "provider.offline-mode.example" +_REMOTE_URLS: Final = { + "LITELLM_MODEL_COST_MAP_URL": "https://cost-map.offline-mode.example/model_prices.json", + "LITELLM_ANTHROPIC_BETA_HEADERS_URL": "https://beta-headers.offline-mode.example/config.json", + "LITELLM_AUTOROUTER_PRESETS_URL": "https://autorouter.offline-mode.example/presets.json", + "LITELLM_BLOG_POSTS_URL": "https://blog.offline-mode.example/rss.xml", +} +_SELF_INITIATED_HOSTS: Final = frozenset( + { + "cost-map.offline-mode.example:443", + "beta-headers.offline-mode.example:443", + "autorouter.offline-mode.example:443", + "blog.offline-mode.example:443", + "raw.githubusercontent.com:443", + "docs.litellm.ai:443", + "api.pwnedpasswords.com:443", + "huggingface.co:443", + } +) +_PER_FEATURE_FLAGS: Final = ( + "LITELLM_OFFLINE_MODE", + "LITELLM_LOCAL_MODEL_COST_MAP", + "LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", + "LITELLM_LOCAL_AUTOROUTER_PRESETS", + "LITELLM_LOCAL_POLICY_TEMPLATES", + "LITELLM_LOCAL_BLOG_POSTS", + "LITELLM_DISABLE_LATEST_RELEASE_CHECK", + "LITELLM_LICENSE", +) + + +def _bundled(name: str) -> JsonValue: + return json.loads((_PACKAGE / name).read_text(encoding="utf-8")) + + +def _write_config(directory: Path) -> Path: + config: Final = directory / f"offline_mode_{uuid.uuid4().hex}.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": "offline-mode-control", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"https://{_CONTROL_HOST}/v1", + "api_key": "synthetic-provider-key", + }, + }, + { + "model_name": "offline-mode-kimi", + "litellm_params": { + "model": "bedrock/us.moonshotai.kimi-k3", + "aws_region_name": "us-east-1", + "aws_access_key_id": "synthetic-access-key", + "aws_secret_access_key": "synthetic-secret-key", + }, + }, + ], + "router_settings": {"num_retries": 0}, + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + } + ) + ) + return config + + +def _json(proxy: Gateway, method: str, path: str, body: dict[str, JsonValue] | None = None) -> JsonValue: + response: Final = proxy.request(method, path, body) + assert response.status_code == 200, f"{method} {path}: {response.status_code} {response.text}" + return json.loads(response.content) + + +@pytest.mark.timeout(300) +@pytest.mark.parametrize( + ("offline_mode", "expected_hosts", "reload_status"), + [ + pytest.param(None, _SELF_INITIATED_HOSTS | {f"{_CONTROL_HOST}:443"}, 502, id="unset"), + pytest.param("true", frozenset({f"{_CONTROL_HOST}:443"}), 200, id="true"), + ], +) +def test_offline_mode_stops_every_self_initiated_request( + gateway: Gateway, tmp_path: Path, offline_mode: str | None, expected_hosts: frozenset[str], reload_status: int +) -> None: + with refusing_forward_proxy() as egress: + environment: Final = { + "HTTPS_PROXY": egress.url, + "HTTP_PROXY": egress.url, + "ALL_PROXY": egress.url, + "NO_PROXY": "127.0.0.1,localhost", + "AIOHTTP_TRUST_ENV": "True", + **_REMOTE_URLS, + **({"LITELLM_OFFLINE_MODE": offline_mode} if offline_mode is not None else {}), + } + with owned_proxy( + gateway, + tmp_path, + environment, + config=_write_config(tmp_path), + remove_environment=_PER_FEATURE_FLAGS, + workers=2, + ) as isolated: + source: Final = object_value(_json(isolated, "GET", "/model/cost_map/source")) + reloaded: Final = isolated.request("POST", "/reload/model_cost_map") + beta_headers: Final = object_value(_json(isolated, "POST", "/reload/anthropic_beta_headers")) + release: Final = _json(isolated, "GET", "/get/latest_release_info") + presets: Final = object_value(_json(isolated, "GET", "/public/autorouter_presets")) + templates: Final = _json(isolated, "GET", "/policy/templates") + blog: Final = object_value(_json(isolated, "GET", "/public/litellm_blog_posts")) + user: Final = object_value( + _json(isolated, "POST", "/user/new", {"user_email": f"offline-{uuid.uuid4().hex}@example.com"}) + ) + password_set: Final = object_value( + _json( + isolated, + "POST", + "/user/update", + {"user_id": user["user_id"], "password": f"Offline-{uuid.uuid4()}"}, + ) + ) + tokens: Final = object_value( + _json( + isolated, + "POST", + "/utils/token_counter", + {"model": "meta-llama/llama-3-8b-instruct", "messages": [{"role": "user", "content": "count me"}]}, + ) + ) + kimi_tokens: Final = object_value( + _json( + isolated, + "POST", + "/utils/token_counter", + {"model": "offline-mode-kimi", "messages": [{"role": "user", "content": "count me"}]}, + ) + ) + control: Final = isolated.request( + "POST", + "/v1/chat/completions", + {"model": "offline-mode-control", "messages": [{"role": "user", "content": "control"}]}, + ) + reached: Final = frozenset(egress.targets()) + + assert reached == expected_hosts + assert control.status_code != 200, control.text + assert source["source"] == "local", source + assert source["is_env_forced"] is (offline_mode is not None), source + assert reloaded.status_code == reload_status, reloaded.text + assert beta_headers["status"] == "success", beta_headers + assert release is None + assert sorted(presets) == sorted(object_value(_bundled("proxy/public_endpoints/autorouter_presets.json"))) + assert templates == _bundled("policy_templates_backup.json") + assert [object_value(post)["url"] for post in _list(blog["posts"])] == [ + object_value(post)["url"] for post in _list(object_value(_bundled("blog_posts.json"))["posts"]) + ] + assert password_set["user_id"] == user["user_id"], password_set + assert tokens["tokenizer_type"] == "openai_tokenizer", tokens + assert isinstance(kimi_tokens["total_tokens"], int) and kimi_tokens["total_tokens"] > 0, kimi_tokens + + +def _list(value: JsonValue) -> list[JsonValue]: + assert isinstance(value, list), value + return value diff --git a/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py b/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py index 780ee43fb3f..0c18f2a4b6f 100644 --- a/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py +++ b/tests/integration/spend/test_daily_activity_key_metadata_query_timeout.py @@ -99,16 +99,9 @@ class _Activity(BaseModel): results: tuple[_Day, ...] -class _UpstreamMessage(BaseModel): - model_config = ConfigDict(frozen=True) - role: str | None = None - - class _UpstreamRequest(BaseModel): model_config = ConfigDict(frozen=True) stream: bool | None = None - tools: tuple[JsonValue, ...] | None = None - messages: tuple[_UpstreamMessage, ...] = () @dataclass(frozen=True, slots=True) @@ -171,17 +164,6 @@ def _stream_frames() -> tuple[bytes, ...]: return (*(f"data: {json.dumps({**envelope, **frame})}\n\n".encode() for frame in frames), b"data: [DONE]\n\n") -def _usage_tool_call() -> Mapping[str, JsonValue]: - arguments: Final = json.dumps({"start_date": _day(-1), "end_date": _day(1)}) - return { - "role": "assistant", - "content": None, - "tool_calls": [ - {"id": "call_usage", "type": "function", "function": {"name": "get_usage_data", "arguments": arguments}} - ], - } - - def _response_object() -> bytes: return json.dumps( { @@ -216,9 +198,6 @@ def _respond(request: Request) -> Reply: body: Final = _UpstreamRequest.model_validate_json(request.body or b"{}") if body.stream: return Reply(content_type="text/event-stream", chunks=_stream_frames()) - roles: Final = frozenset(message.role for message in body.messages) - if body.tools and "tool" not in roles: - return Reply(body=_completion(_usage_tool_call(), "tool_calls")) return Reply(body=_completion({"role": "assistant", "content": REPLY_TEXT}, "stop")) @@ -459,12 +438,6 @@ def _restore(database_url: str, digest: str) -> None: ) -def _events(response: httpx.Response) -> tuple[Mapping[str, JsonValue], ...]: - return tuple( - json.loads(line.removeprefix("data: ")) for line in response.text.splitlines() if line.startswith("data: ") - ) - - @dataclass(frozen=True, slots=True) class Tenant: label: str @@ -714,39 +687,19 @@ def _rig(gateway: Gateway, directory: Path) -> Generator[tuple[Gateway, str]]: @pytest.mark.timeout(360) -def test_usage_ai_chat_timeout_then_second_timeout_still_recovers_key_alias_after_database_frees( +def test_usage_page_timeout_then_second_timeout_still_recovers_key_alias_after_database_frees( gateway: Gateway, tmp_path: Path ) -> None: with _rig(gateway, tmp_path) as (proxy, database_url): - spender: Final = _spender(proxy, database_url, "ai-chat") + spender: Final = _spender(proxy, database_url, "repeated-timeout") with _pinned(proxy) as pinned: with _locked_spend_logs(database_url): - with _recording(database_url) as during_chat: - chat: Final = pinned.request( - "POST", - "/usage/ai/chat", - body={ - "messages": [{"role": "user", "content": "What did we spend?"}], - "model": "openai/gpt-4o-mini", - }, - ) - assert chat.status_code == 200, chat.text - assert chat.elapsed >= FAILED_LOOKUP_FLOOR, chat.elapsed - tool_call: Final = { - "type": "tool_call", - "tool_name": "get_usage_data", - "tool_label": "global usage data", - "arguments": {"start_date": _day(-1), "end_date": _day(1)}, - } - assert _events(chat) == ( - {"type": "status", "message": "Thinking..."}, - {**tool_call, "status": "running"}, - {**tool_call, "status": "complete"}, - {"type": "status", "message": "Analyzing results..."}, - {"type": "chunk", "content": REPLY_TEXT}, - {"type": "done"}, - ), chat.text - assert len(_lookups(during_chat)) == 1 + with _recording(database_url) as during_timeout: + timed_out: Final = _read(pinned, spender.digest) + assert timed_out.status_code == 200, timed_out.text + assert timed_out.elapsed >= FAILED_LOOKUP_FLOOR, timed_out.elapsed + assert _aliases(timed_out, spender.digest) == (None,), timed_out.text + assert len(_lookups(during_timeout)) == 1 cached_miss: Final = _read(pinned, spender.digest) assert cached_miss.elapsed < FAILED_LOOKUP_FLOOR, cached_miss.elapsed assert _aliases(cached_miss, spender.digest) == (None,), cached_miss.text @@ -1087,44 +1040,3 @@ def test_usage_page_rejects_bad_key_filters_and_unrelated_routes_ignore_a_locked assert gateway_activity == 200, gateway_activity chat: Final = proxy.chat(MODEL, text=f"locked {uuid.uuid4().hex}") assert object_value(chat["usage"]) == dict(USAGE), chat - - -@pytest.mark.timeout(300) -def test_usage_ai_chat_survives_a_dropped_database_connection_during_alias_recovery( - gateway: Gateway, tmp_path: Path -) -> None: - with ( - scratch_database() as database_url, - database_relay(database_url, b"AS " + LOOKUP_MARKER.encode()) as (relay, relayed_url), - wire_server(_respond) as wire, - _proxy(gateway, tmp_path, relayed_url, wire.url) as proxy, - ): - spender: Final = _spender(proxy, database_url, "dropped-ai-chat") - with _pinned(proxy) as pinned: - relay.arm() - chat: Final = pinned.request( - "POST", - "/usage/ai/chat", - body={ - "messages": [{"role": "user", "content": "What did we spend?"}], - "model": "openai/gpt-4o-mini", - }, - ) - assert chat.status_code == 200, chat.text - assert relay.tripped.is_set(), chat.text - tool_call: Final = { - "type": "tool_call", - "tool_name": "get_usage_data", - "tool_label": "global usage data", - "arguments": {"start_date": _day(-1), "end_date": _day(1)}, - } - assert _events(chat) == ( - {"type": "status", "message": "Thinking..."}, - {**tool_call, "status": "running"}, - {**tool_call, "status": "complete"}, - {"type": "status", "message": "Analyzing results..."}, - {"type": "chunk", "content": REPLY_TEXT}, - {"type": "done"}, - ), chat.text - recovered: Final = _named_once_reconnected(pinned, spender) - assert _metadata(recovered, spender.digest) == _full_metadata(spender), recovered.text diff --git a/tests/integration/spend/test_daily_activity_key_owner_traffic.py b/tests/integration/spend/test_daily_activity_key_owner_traffic.py index b8f113aca49..5284ad631de 100644 --- a/tests/integration/spend/test_daily_activity_key_owner_traffic.py +++ b/tests/integration/spend/test_daily_activity_key_owner_traffic.py @@ -6,7 +6,6 @@ from collections.abc import Iterable from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime, timedelta from hashlib import sha256 -from pathlib import Path from queue import SimpleQueue from typing import Final @@ -31,8 +30,7 @@ from integration._support.daily_activity import ( user_row, user_with_an_email, ) -from integration._support.database import read_rows, scratch_database -from integration._support.process import owned_proxy +from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -51,22 +49,6 @@ UNIFIED_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/respons REQUESTS_OF_A_BURST: Final = 21 READS_DURING_A_BURST: Final = 30 TOKEN_LIMIT_DISCOVERY: Final = ("GET", "/v1/models") -TOOL_CALL: Final = "call_integration_usage" -ANSWER: Final = "One request cost $0.25" -SUMMARY_OF_ONE_SEEDED_ROW: Final = "\n".join( - ( - "Total Spend: $0.2500", - "Total Requests: 1", - "Successful: 1 | Failed: 0", - "Total Tokens: 15", - "", - "Top Models by Spend:", - " - gpt-4o-mini: $0.2500 (1 reqs, 15 tokens)", - "", - "Top Providers by Spend:", - " - openai: $0.2500 (1 reqs)", - ) -) def _chat_completion() -> dict[str, JsonValue]: @@ -105,55 +87,6 @@ def _provider(request: Request) -> Reply: return Reply(body=json.dumps(body).encode()) -def _usage_tool_call() -> dict[str, JsonValue]: - call: Final[dict[str, JsonValue]] = { - "id": TOOL_CALL, - "type": "function", - "function": { - "name": "get_usage_data", - "arguments": json.dumps({"start_date": DAY, "end_date": DAY}), - }, - } - return { - "id": f"chatcmpl-{uuid.uuid4().hex}", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": None, "tool_calls": [call]}, - "finish_reason": "tool_calls", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - } - - -def _streamed_chunk(delta: dict[str, JsonValue], finish_reason: str | None) -> bytes: - chunk: Final = { - "id": "chatcmpl-integration-usage", - "object": "chat.completion.chunk", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], - } - return f"data: {json.dumps(chunk)}\n\n".encode() - - -def _usage_analyst(request: Request) -> Reply: - if json.loads(request.body).get("stream"): - return Reply( - chunks=( - _streamed_chunk({"role": "assistant", "content": ANSWER}, None), - _streamed_chunk({}, "stop"), - b"data: [DONE]\n\n", - ), - content_type="text/event-stream", - ) - return Reply(body=json.dumps(_usage_tool_call()).encode()) - - def _sent_for_callers(requests: Iterable[Request]) -> tuple[Request, ...]: return tuple(request for request in requests if (request.method, request.target) != TOKEN_LIMIT_DISCOVERY) @@ -322,65 +255,6 @@ def test_cli_session_spend_is_reported_with_the_user_and_team_of_the_session( ) -@pytest.mark.timeout(300) -def test_usage_ai_chat_hands_the_model_the_usage_summary_without_any_key_owner( - gateway: Gateway, tmp_path: Path -) -> None: - question: Final = f"what did we spend {uuid.uuid4().hex}" - owner: Final = f"integration-{uuid.uuid4().hex}" - ownerless_key: Final = f"integration-ownerless-{uuid.uuid4().hex}" - with ( - scratch_database() as scratch_url, - wire_server(_usage_analyst) as wire, - owned_proxy( - gateway, - tmp_path, - { - "DATABASE_URL": scratch_url, - "OPENAI_API_BASE": f"{wire.url}/v1", - "OPENAI_BASE_URL": f"{wire.url}/v1", - "OPENAI_API_KEY": "integration-provider-key", - }, - remove_environment=("DATABASE_URL_READ_REPLICA",), - ) as candidate, - ): - candidate.post("/user/new", {"user_id": owner, "user_email": f"{owner}@example.com", "auto_create_key": False}) - with daily_rows((user_row(owner, ownerless_key, DAY),), database_url=scratch_url): - answer: Final = candidate.request( - "POST", - "/usage/ai/chat", - {"messages": [{"role": "user", "content": question}], "model": "openai/gpt-4o-mini"}, - ) - assert answer.status_code == 200, answer.text - tool_call: Final = { - "type": "tool_call", - "tool_name": "get_usage_data", - "tool_label": "global usage data", - "arguments": {"start_date": DAY, "end_date": DAY}, - } - events: Final = [ - json.loads(line.removeprefix("data: ")) for line in answer.text.splitlines() if line.startswith("data: ") - ] - assert events == [ - {"type": "status", "message": "Thinking..."}, - {**tool_call, "status": "running"}, - {**tool_call, "status": "complete"}, - {"type": "status", "message": "Analyzing results..."}, - {"type": "chunk", "content": ANSWER}, - {"type": "done"}, - ], answer.text - asked, analysed = wire.drain() - assert [asked.target, analysed.target] == ["/v1/chat/completions", "/v1/chat/completions"] - assert json.loads(asked.body)["messages"][-1] == {"role": "user", "content": question} - assert json.loads(analysed.body)["messages"][-1] == { - "role": "tool", - "tool_call_id": TOOL_CALL, - "content": SUMMARY_OF_ONE_SEEDED_ROW, - } - assert owner not in analysed.body.decode() - assert ownerless_key not in analysed.body.decode() - - @pytest.mark.timeout(300) def test_owner_is_reported_on_every_route_while_a_burst_of_requests_waits_on_the_provider(gateway: Gateway) -> None: released: Final = threading.Event() diff --git a/tests/integration/spend/test_spend_log_tool_payload_content.py b/tests/integration/spend/test_spend_log_tool_payload_content.py index 9055f6f06c4..226accdbf54 100644 --- a/tests/integration/spend/test_spend_log_tool_payload_content.py +++ b/tests/integration/spend/test_spend_log_tool_payload_content.py @@ -15,7 +15,7 @@ import openai 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 owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -723,6 +723,33 @@ def test_chat_sdk_stream_include_usage_masks_request_fields(gateway: Gateway, tm assert len(post_requests) == 1 +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_chat_stream_usage_the_proxy_adds_keeps_its_marker_out_of_the_stored_request( + gateway: Gateway, tmp_path: Path +) -> None: + response_id: Final = f"chatcmpl-proxy-usage-{uuid4()}" + with ( + wire_server( + _answering_model_listing(lambda request: _chat_stream(response_id, "sort_key", include_usage=True)) + ) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, + isolated.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") + response: Final = isolated.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "stream marker"}], "stream": True}, + ) + assert response.status_code == 200, response.text + assert {string_value(event["id"]) for event in _sse_events(response.text)} == {response_id}, response.text + stored_request: Final = object_value(_stored_row(response_id)["proxy_server_request"]) + assert stored_request["messages"] == [{"role": "user", "content": "stream marker"}], stored_request + assert "_litellm_stream_usage_injected" not in stored_request, stored_request + observed: Final = _provider_calls(wire.drain()) + assert [_json_object(request.body)["stream_options"] for request in observed] == [{"include_usage": True}] + + def test_chat_history_keeps_string_tool_arguments_and_tool_content(gateway: Gateway, tmp_path: Path) -> None: arguments: Final = '{"sort_key":"created_at"}' tool_content: Final = "tool-result-created_at" diff --git a/tests/integration/spend/test_team_member_total_spend_reconciliation.py b/tests/integration/spend/test_team_member_total_spend_reconciliation.py new file mode 100644 index 00000000000..f5ccae0d12b --- /dev/null +++ b/tests/integration/spend/test_team_member_total_spend_reconciliation.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import asyncio +import json +import os +import uuid +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import httpx +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob +from litellm.proxy.utils import PrismaClient, ProxyLogging +from tests.integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value +from tests.integration._support.database import read_rows, write_rows +from tests.integration._support.provider import PROVIDER_URL, SharedProvider +from tests.integration._support.wire import Reply + +_CALL_COST: Final = 0.02 +JSON_VALUE: Final = TypeAdapter(JsonValue) +_LEGACY_SPEND: Final = 612.4 +_LEGACY_TOTAL: Final = 148.1 +_HEALTHY_SPEND: Final = 12.25 +_HEALTHY_TOTAL: Final = 97.8 + + +def _completion() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + +def _priced_model(scenario: Scenario) -> str: + return scenario.model(api_base=f"{PROVIDER_URL}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002) + + +def _member_with_budget(gateway: Gateway, scenario: Scenario, team_id: str, *, max_budget: float) -> str: + user_id: Final = scenario.user(user_role="internal_user") + gateway.post( + "/team/member_add", + {"team_id": team_id, "member": {"role": "user", "user_id": user_id}, "max_budget_in_team": max_budget}, + ) + return user_id + + +def _seed_membership(user_id: str, team_id: str, *, spend: float, total_spend: float) -> None: + write_rows( + 'UPDATE "LiteLLM_TeamMembership" SET spend = %s::float8, total_spend = %s::float8 ' + "WHERE user_id = %s AND team_id = %s", + (str(spend), str(total_spend), user_id, team_id), + ) + + +def _membership(user_id: str, team_id: str) -> tuple[float, float]: + rows: Final = read_rows( + 'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s', + (user_id, team_id), + ) + assert len(rows) == 1, rows + return float(str(rows[0]["spend"])), float(str(rows[0]["total_spend"])) + + +def _membership_reaches(user_id: str, team_id: str, *, total_spend: float) -> tuple[float, float]: + return eventually( + lambda: _membership(user_id, team_id), + lambda row: row[1] == pytest.approx(total_spend), + seconds=70, + ) + + +def _served(gateway: Gateway, provider: SharedProvider, model: str, key: str) -> None: + provider.expect(_completion()) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"reconcile {uuid.uuid4().hex}"}]}, + key=key, + ) + assert response.status_code == 200, response.text + assert len(provider.received()) == 1 + + +def _reset_spend(gateway: Gateway, team_id: str, user_id: str, reset_to: float) -> httpx.Response: + return gateway.request("POST", f"/team/{team_id}/member/{user_id}/reset_spend", {"reset_to": reset_to}) + + +def _membership_figures(row: JsonValue) -> tuple[str, tuple[float, float]]: + membership: Final = object_value(row) + return string_value(membership["user_id"]), ( + float(str(membership["spend"])), + float(str(membership["total_spend"])), + ) + + +def _reported_memberships(gateway: Gateway, path: str, params: dict[str, str]) -> dict[str, tuple[float, float]]: + """/team/info answers with one team object, /team/list with a list of them.""" + response: Final = gateway.request("GET", path, params=params) + assert response.status_code == 200, f"GET {path}: {response.status_code} {response.text}" + body: Final = JSON_VALUE.validate_json(response.content) + teams: Final = [object_value(body)] if isinstance(body, dict) else [object_value(team) for team in list_value(body)] + rows: Final = [ + row for team in teams for row in list_value(team["team_memberships"]) + ] # comprehension-ok: flatten memberships across teams + return dict(_membership_figures(row) for row in rows) + + +def test_a_flush_lifts_a_legacy_total_to_the_cycle_spend_and_leaves_a_healthy_row_on_its_baseline( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + team_id: Final = scenario.team(models=[model]) + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + _seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL) + legacy_key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model]) + healthy_key: Final = scenario.key(user_id=healthy, team_id=team_id, models=[model]) + + _served(gateway, provider, model, legacy_key) + _served(gateway, provider, model, healthy_key) + + legacy_row: Final = _membership_reaches( + legacy, team_id, total_spend=max(_LEGACY_SPEND, _LEGACY_TOTAL) + _CALL_COST + ) + assert legacy_row[0] == pytest.approx(_LEGACY_SPEND + _CALL_COST) + healthy_row: Final = _membership_reaches(healthy, team_id, total_spend=_HEALTHY_TOTAL + _CALL_COST) + assert healthy_row[0] == pytest.approx(_HEALTHY_SPEND + _CALL_COST) + + +def test_team_info_and_team_list_report_a_legacy_total_no_lower_than_the_cycle_spend(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + _seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL) + + info: Final = _reported_memberships(gateway, "/team/info", {"team_id": team_id}) + listed: Final = _reported_memberships(gateway, "/team/list", {"team_id": team_id}) + + for reported in (info, listed): + assert reported[legacy] == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_SPEND)) + assert reported[healthy] == (pytest.approx(_HEALTHY_SPEND), pytest.approx(_HEALTHY_TOTAL)) + assert _membership(legacy, team_id) == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_TOTAL)) + + +@pytest.mark.parametrize("reset_to", [0.0, 100.0], ids=["to_zero", "to_nonzero"]) +def test_an_admin_reset_keeps_the_legacy_cycle_spend_as_the_lifetime_floor(gateway: Gateway, reset_to: float) -> None: + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + + response: Final = _reset_spend(gateway, team_id, legacy, reset_to) + + assert response.status_code == 200, response.text + assert _membership(legacy, team_id) == (pytest.approx(reset_to), pytest.approx(_LEGACY_SPEND)) + + +def test_concurrent_admin_resets_leave_one_winner_and_a_lifetime_floor_at_the_old_cycle_spend( + gateway: Gateway, +) -> None: + """Two admins reset the same member at once. Each UPDATE re-checks reset_to <= spend, so + the higher reset either lands first or is refused with 409 once the lower one has landed. + Whatever the order, the row ends at the lower value with the old cycle spend preserved.""" + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + + with ThreadPoolExecutor(max_workers=2) as pool: + higher, lower = pool.map(lambda value: _reset_spend(gateway, team_id, legacy, value), (600.0, 500.0)) + + assert lower.status_code == 200, lower.text + assert higher.status_code in (200, 409), higher.text + assert _membership(legacy, team_id) == (pytest.approx(500.0), pytest.approx(_LEGACY_SPEND)) + + +def test_an_admin_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + team_id: Final = scenario.team(models=[model]) + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model]) + + _served(gateway, provider, model, key) + reset: Final = _reset_spend(gateway, team_id, legacy, 0.0) + assert reset.status_code == 200, reset.text + + spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST) + assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total) + assert total >= spend + + +def _expire_member_budget(user_id: str, team_id: str, *, budget_duration: str) -> None: + write_rows( + "UPDATE \"LiteLLM_BudgetTable\" SET budget_duration = %s, budget_reset_at = NOW() - INTERVAL '1 hour' " + 'WHERE budget_id = (SELECT budget_id FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s)', + (budget_duration, user_id, team_id), + ) + + +async def _run_scheduled_reset() -> None: + client: Final = PrismaClient( + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()) + ) + await client.connect() + try: + await ResetBudgetJob( + ProxyLogging(user_api_key_cache=DualCache()), client + ).reset_budget_for_litellm_budget_table() + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize( + "spend, max_budget, carried", + [(8.0, 10.0, 0.0), (10.0, 10.0, 0.0), (15.0, 10.0, 5.0)], + ids=["below_cap", "at_cap", "above_cap_carries_overage"], +) +def test_a_scheduled_reset_lifts_the_legacy_total_before_zeroing_or_carrying_the_cycle_spend( + gateway: Gateway, monkeypatch: pytest.MonkeyPatch, spend: float, max_budget: float, carried: float +) -> None: + monkeypatch.setattr(litellm, "budget_rollover", True) + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=max_budget) + _seed_membership(legacy, team_id, spend=spend, total_spend=3.0) + _expire_member_budget(legacy, team_id, budget_duration="1d") + + asyncio.run(_run_scheduled_reset()) + + assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend)) + asyncio.run(_run_scheduled_reset()) + assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend)) + + +def test_a_scheduled_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + team_id: Final = scenario.team(models=[model]) + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + _expire_member_budget(legacy, team_id, budget_duration="1d") + key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model]) + + _served(gateway, provider, model, key) + asyncio.run(_run_scheduled_reset()) + + spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST) + assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total) + assert total >= spend diff --git a/tests/integration/streaming/test_stream_contracts.py b/tests/integration/streaming/test_stream_contracts.py index d7537a78db8..1d2a8b25792 100644 --- a/tests/integration/streaming/test_stream_contracts.py +++ b/tests/integration/streaming/test_stream_contracts.py @@ -3,6 +3,7 @@ import json import threading import uuid from collections.abc import Callable +from itertools import chain from pathlib import Path from queue import SimpleQueue from typing import Final @@ -13,9 +14,10 @@ from hypothesis import Phase, example, given, settings from hypothesis import strategies as st from integration._support.client import Gateway, 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 -from openai import OpenAI +from openai import OpenAI, omit +from openai.types.chat import ChatCompletionChunk from pydantic import JsonValue, TypeAdapter pytestmark: Final = pytest.mark.shares_worker_state @@ -197,6 +199,7 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat assert {chunk.id for chunk in chunks} == {identity} usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) assert len(usages) == (1 if include else 0) + assert [len(chunk.choices) for chunk in chunks if chunk.usage is not None] == ([0] if include else []) if include: assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4 requests: Final = wire.drain() @@ -214,6 +217,119 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat assert float(rows[0]["spend"]) == pytest.approx(0.019) +_STREAM_USAGE_INJECTED_FLAG: Final = "_litellm_stream_usage_injected" + + +def _registered_model_names(gateway: Gateway) -> tuple[JsonValue, ...]: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list), entries + return tuple(object_value(entry).get("model_name") for entry in entries) + + +def _registered(gateway: Gateway, model: str) -> None: + eventually(lambda: _registered_model_names(gateway), lambda model_names: model in model_names, seconds=60) + + +def _sdk_stream( + gateway: Gateway, + route_prefix: str, + model: str, + identity: str, + *, + include_usage: bool | None, + extra_body: dict[str, JsonValue], +) -> tuple[ChatCompletionChunk, ...]: + with OpenAI( + api_key=gateway.key, + base_url=str(gateway.client.base_url).rstrip("/") + route_prefix, + timeout=30, + max_retries=0, + ) as client: + stream: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": identity}], + stream=True, + stream_options=omit if include_usage is None else {"include_usage": include_usage}, + extra_body=extra_body, + ) + with stream: + return tuple(stream) + + +def _streamed_text(chunks: tuple[ChatCompletionChunk, ...]) -> str: + choices: Final = chain.from_iterable(chunk.choices for chunk in chunks) + return "".join(choice.delta.content or "" for choice in choices) + + +def _usage_chunk_shapes(chunks: tuple[ChatCompletionChunk, ...]) -> list[tuple[int, int]]: + return [(len(chunk.choices), chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None] + + +@pytest.mark.parametrize( + ("route_prefix", "sends", "extra_body"), + [pytest.param("/v1", 2, {}, id="chat_miss_then_cache_hit"), pytest.param("/queue", 1, {"priority": 0}, id="queue")], +) +def test_caller_requested_stream_usage_has_no_choices_even_with_a_client_sent_injection_flag( + gateway: Gateway, route_prefix: str, sends: int, extra_body: dict[str, JsonValue] +) -> None: + identity: Final = "stream-usage-flag-" + uuid.uuid4().hex + with ( + wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=wire.url + "/v1") + _registered(gateway, model) + streams: Final = tuple( + _sdk_stream( + gateway, + route_prefix, + model, + identity, + include_usage=True, + extra_body={**extra_body, _STREAM_USAGE_INJECTED_FLAG: True}, + ) + for _ in range(sends) + ) + assert [_streamed_text(chunks) for chunks in streams] == ["Hello 雪 café"] * sends + assert [_usage_chunk_shapes(chunks) for chunks in streams] == [[(0, 15)]] * sends + requests: Final = wire.drain() + assert len(requests) == 1 + assert _STREAM_USAGE_INJECTED_FLAG not in json.loads(requests[0].body) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_stream_usage_the_deployment_adds_keeps_one_empty_choice(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "deployment-stream-usage-" + uuid.uuid4().hex + with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire: + shared_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **shared_config, + "model_list": [ + { + "model_name": "deployment-usage", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "synthetic-deployment-usage-key", + "api_base": wire.url + "/v1", + "stream_options": {"include_usage": True}, + }, + } + ], + "general_settings": {**shared_config["general_settings"], "always_include_stream_usage": False}, + } + path: Final = tmp_path / "deployment-stream-usage.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(gateway, tmp_path, {}, config=path) as candidate: + chunks: Final = _sdk_stream( + candidate, "/v1", "deployment-usage", identity, include_usage=None, extra_body={} + ) + assert _streamed_text(chunks) == "Hello 雪 café" + assert _usage_chunk_shapes(chunks) == [(1, 15)] + requests: Final = wire.drain() + assert len(requests) == 1 + assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True + + @pytest.mark.covers("other.streaming.messages_bridge.empty_choices_usage_chunk_completes_stream") def test_messages_stream_completes_through_trailing_empty_choices_usage_chunk(gateway: Gateway) -> None: identity: Final = "messages-empty-choices-" + uuid.uuid4().hex diff --git a/tests/unit/AGENTS.md b/tests/unit/AGENTS.md index 9ce6c58cff7..ca717e8efb0 100644 --- a/tests/unit/AGENTS.md +++ b/tests/unit/AGENTS.md @@ -1,6 +1,9 @@ # tests/unit -In-process. No network or subprocess +In-process. No network or subprocess. `interactions/test_openapi_compliance.py` reads the spec snapshot beside it, +and only the `interactions-spec-canary` workflow passes `--interactions-spec-url` (registered in +`interactions/conftest.py`, so name that file or directory on the command line) with `ai.google.dev` allowed, to run +the same tests against Google's live spec Unit tests must not depend on the current time or execution speed. Use fixed timestamps for fixtures and controlled clocks for time-dependent behavior. Incidental clock reads inside production code are allowed diff --git a/tests/unit/a2a_protocol/test_utils.py b/tests/unit/a2a_protocol/test_utils.py index 8a219211c62..dfa8449f6e6 100644 --- a/tests/unit/a2a_protocol/test_utils.py +++ b/tests/unit/a2a_protocol/test_utils.py @@ -1,5 +1,8 @@ """Tests for litellm/a2a_protocol/utils.py token/usage extraction.""" +from types import SimpleNamespace +from typing import Final + import pytest pytest.importorskip("a2a.compat.v0_3.types") @@ -7,6 +10,7 @@ pytest.importorskip("a2a.compat.v0_3.types") from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest from litellm.a2a_protocol.utils import A2ARequestUtils +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder def _request(user_text: str) -> SendMessageRequest: @@ -39,3 +43,13 @@ def test_calculate_usage_counts_input_tokens_from_request_object(): assert prompt_tokens > 0 assert completion_tokens > 0 assert total_tokens == prompt_tokens + completion_tokens + + +def test_calculate_usage_counts_input_tokens_from_a_message_that_forwards_its_fields(): + request: Final = _request("count these input tokens please") + forwarded_request: Final = SimpleNamespace(params=SimpleNamespace(message=AttributeForwarder(request.params.message))) + response_dict: Final = {"result": {"kind": "message", "parts": [{"kind": "text", "text": "ok"}]}} + + assert A2ARequestUtils.calculate_usage_from_request_response( + request=forwarded_request, response_dict=response_dict + ) == A2ARequestUtils.calculate_usage_from_request_response(request=request, response_dict=response_dict) diff --git a/tests/unit/caching/test_in_memory_cache.py b/tests/unit/caching/test_in_memory_cache.py index 214f4ecf301..75232476ecd 100644 --- a/tests/unit/caching/test_in_memory_cache.py +++ b/tests/unit/caching/test_in_memory_cache.py @@ -49,6 +49,47 @@ async def test_async_increment_delegates_to_locked_sync_path(): assert cache.get_cache("counter") == 5 +class _AnotherReaderExpiresTheEntryFirstCache(InMemoryCache): + other_reader_started = False + other_reader_result: object = "not run" + + def evict_element_if_expired(self, key: str) -> bool: + if not self.other_reader_started: + self.other_reader_started = True + self.other_reader_result = self.get_cache(key) + return super().evict_element_if_expired(key) + + +def test_get_cache_reads_an_entry_another_reader_just_expired_as_a_miss() -> None: + cache: Final = _AnotherReaderExpiresTheEntryFirstCache() + cache.set_cache("tokenizer_failure", "HF unavailable", ttl=-1) + + assert cache.get_cache("tokenizer_failure") is None + assert cache.other_reader_result is None + + +class _ClockThatLetsAnotherReaderRunFirst: + cache: InMemoryCache | None = None + other_reader_result: object = "not run" + + def __call__(self) -> float: + cache: Final = self.cache + if cache is not None: + self.cache = None + self.other_reader_result = cache.get_cache("tokenizer_failure") + return 100.0 + + +def test_get_cache_reads_an_entry_another_reader_expires_during_the_expiry_check_as_a_miss() -> None: + clock: Final = _ClockThatLetsAnotherReaderRunFirst() + cache: Final = InMemoryCache(clock=clock) + cache.set_cache("tokenizer_failure", "HF unavailable", ttl=-1) + clock.cache = cache + + assert cache.get_cache("tokenizer_failure") is None + assert clock.other_reader_result is None + + def test_in_memory_openai_obj_cache(): from openai import OpenAI diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index c96e3e2c031..e64f5cca5b8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -252,6 +252,7 @@ def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> It @pytest.fixture(autouse=True) def isolate_litellm_globals() -> Iterator[None]: original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS} + original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks) original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS} original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)} original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__} @@ -263,6 +264,7 @@ def isolate_litellm_globals() -> Iterator[None]: _flush_client_caches() for name in CALLBACK_LISTS: setattr(litellm, name, []) + litellm_router_module._routers_holding_global_callbacks.clear() for name in RESET_TO_NONE_GLOBALS: setattr(litellm, name, None) for name in MODULE_LEVEL_CLIENTS: @@ -277,6 +279,8 @@ def isolate_litellm_globals() -> Iterator[None]: litellm.__dict__.update(original_clients) for name, value in (original_callbacks | original_reset | original_restored).items(): setattr(litellm, name, value) + litellm_router_module._routers_holding_global_callbacks.clear() + litellm_router_module._routers_holding_global_callbacks.update(original_callback_holders) for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items(): logger.setLevel(level) logger.disabled = disabled diff --git a/tests/unit/containers/test_container_proxy_ownership.py b/tests/unit/containers/test_container_proxy_ownership.py index 780ae9bf6d1..ac5adeee8c1 100644 --- a/tests/unit/containers/test_container_proxy_ownership.py +++ b/tests/unit/containers/test_container_proxy_ownership.py @@ -1,6 +1,7 @@ import json import sys from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -10,6 +11,7 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.container_endpoints import ownership from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.containers.main import ContainerListResponse, ContainerObject +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder @pytest.fixture(autouse=True) @@ -97,6 +99,23 @@ async def test_should_not_mutate_dict_container_response_when_recording_owner( assert file_obj["provider_container_id"] == "cntr_provider" +@pytest.mark.asyncio +async def test_should_store_the_full_container_of_a_response_that_forwards_its_fields(monkeypatch): + table: Final = AsyncMock() + table.find_unique.return_value = None + prisma_client: Final = SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table)) + monkeypatch.setattr(ownership, "_get_prisma_client", AsyncMock(return_value=prisma_client)) + + await ownership.record_container_owner( + response=AttributeForwarder(_container("cntr_provider")), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + custom_llm_provider="openai", + ) + + file_obj: Final = json.loads(table.create.await_args.kwargs["data"]["file_object"]) + assert file_obj["status"] == "active" + + @pytest.mark.asyncio async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch): table = AsyncMock() diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 84e7a970195..e2385a2e215 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -3799,3 +3799,35 @@ def test_nace_upstream_failures_preserve_status_and_message( ) assert getattr(error.value, "status_code") == status assert upstream.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + "usage", [None, {"input_tokens": 0, "output_tokens": 0}, {"input_tokens": 12, "output_tokens": 1}] +) +async def test_native_decisions_distinguishes_missing_usage_from_reported_zero( + is_async: bool, usage: dict[str, int] | None, respx_mock: respx.MockRouter +) -> None: + body = { + **{key: value for key, value in _RESPONSE.items() if key != "usage"}, + "model": "jev-1.13.0", + **({"usage": usage} if usage is not None else {}), + } + respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=body) + params = {"model": "typesafe/jev-1.13.0", "state": "review", "questions": _QUESTIONS, "api_key": "synthetic-key"} + response = await litellm.adecisions(**params) if is_async else litellm.decisions(**params) + + assert isinstance(response, DecisionsResponse) + if usage is None: + assert response.usage is None + assert response._hidden_params.get("response_cost") is None + with pytest.raises(ValueError, match="provider did not report usage"): + litellm.completion_cost(completion_response=response) + else: + assert response.usage is not None + assert response.usage.input_tokens == usage["input_tokens"] + assert response.usage.output_tokens == usage["output_tokens"] + expected = litellm.completion_cost(completion_response=response) + assert response._hidden_params["response_cost"] == pytest.approx(expected) + assert (expected == 0) == (usage["input_tokens"] == usage["output_tokens"] == 0) diff --git a/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py index 953f028af3c..cf53fe90216 100644 --- a/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py +++ b/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py @@ -11,14 +11,14 @@ from litellm.endpoints.speech.speech_to_completion_bridge.transformation import ) from litellm.types.utils import ChatCompletionAudioResponse, Choices, Message, ModelResponse +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + GEMINI_TTS_MODEL: Final = "gemini-3.1-flash-tts-preview" PCM_BYTES: Final = b"\x01\x02\x03\x04" * 6 def _model_response(model: str, pcm: bytes) -> ModelResponse: - audio: Final = ChatCompletionAudioResponse( - data=base64.b64encode(pcm).decode(), expires_at=0, transcript="hello" - ) + audio: Final = ChatCompletionAudioResponse(data=base64.b64encode(pcm).decode(), expires_at=0, transcript="hello") return ModelResponse(model=model, choices=[Choices(message=Message(content=None, audio=audio))]) @@ -72,6 +72,103 @@ def test_non_gemini_request_forwards_speech_response_format_as_audio_format() -> assert request["audio"] == {"voice": "alloy", "format": "wav"} +def test_gemini_request_keeps_voice_beside_nested_speech_config() -> None: + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Bonjour", + voice={"speech_config": {"language_code": "fr-FR"}, "voice": "Kore"}, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"]["speech_config"] == { + "languageCode": "fr-FR", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}, + } + + +def test_gemini_request_keeps_voice_object_for_a_model_alias_map_alias(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"voice-alias": GEMINI_TTS_MODEL}) + + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model="voice-alias", + input="Hello there", + voice={"voice": "Puck"}, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"] == { + "speech_config": {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}}, + "format": "pcm16", + } + + +@pytest.mark.parametrize("wrapper", (None, "speech_config", "speechConfig")) +def test_gemini_request_keeps_voice_and_language_without_audio_format(wrapper: str | None) -> None: + voice: Final = {"language_code": "fr-FR", "voice": "Kore", "format": "pcm16"} + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Bonjour", + voice=voice if wrapper is None else {wrapper: voice}, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"]["speech_config"] == { + "languageCode": "fr-FR", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}, + } + + +@pytest.mark.parametrize( + ("drop_params", "voice_config"), + [ + (True, {"prebuiltVoiceConfig": {"voiceName": "Kore"}}), + (False, {"prebuiltVoiceConfig": {"voiceName": "Kore", "style": "calm"}, "tone": "warm"}), + ], +) +def test_gemini_request_drops_unknown_nested_voice_fields_only_with_drop_params( + drop_params: bool, voice_config: dict[str, object] +) -> None: + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Bonjour", + voice={"voice_config": {"prebuilt_voice_config": {"voice_name": "Kore", "style": "calm"}, "tone": "warm"}}, + optional_params={}, + litellm_params={"drop_params": drop_params}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"]["speech_config"] == {"voiceConfig": voice_config} + + +def test_gemini_request_reads_snake_case_wrapper_beside_null_camel_case_wrapper() -> None: + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Bonjour", + voice={"speechConfig": None, "speech_config": {"voice_name": "Kore"}}, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"]["speech_config"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + @pytest.mark.parametrize("response_format", ["mp3", "flac", "opus", "aac"]) def test_gemini_tts_request_rejects_formats_gemini_cannot_produce(response_format: str) -> None: with pytest.raises(litellm.BadRequestError) as excinfo: @@ -86,6 +183,7 @@ def test_gemini_tts_request_rejects_formats_gemini_cannot_produce(response_forma def test_gemini_tts_pcm_response_returns_raw_pcm_bytes() -> None: response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( model_response=_model_response(GEMINI_TTS_MODEL, PCM_BYTES), + model=GEMINI_TTS_MODEL, response_format="pcm", ) @@ -97,6 +195,7 @@ def test_gemini_tts_pcm_response_returns_raw_pcm_bytes() -> None: def test_gemini_tts_wav_and_default_responses_wrap_pcm_in_wav(response_format: str | None) -> None: response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( model_response=_model_response(GEMINI_TTS_MODEL, PCM_BYTES), + model=GEMINI_TTS_MODEL, response_format=response_format, ) @@ -110,8 +209,24 @@ def test_gemini_tts_wav_and_default_responses_wrap_pcm_in_wav(response_format: s def test_non_gemini_response_keeps_original_bytes_and_mpeg_content_type() -> None: response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( model_response=_model_response("gpt-4o-audio-preview", PCM_BYTES), + model="gpt-4o-audio-preview", response_format="mp3", ) assert response.response.content == PCM_BYTES assert response.response.headers["content-type"] == "audio/mpeg" + + +@pytest.mark.parametrize( + ("requested_model", "served_model"), + [(GEMINI_TTS_MODEL, "served-speech-alias"), ("requested-speech-alias", GEMINI_TTS_MODEL)], +) +def test_gemini_response_wraps_wav_when_either_model_name_is_gemini_tts(requested_model: str, served_model: str) -> None: + response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( + model_response=_model_response(served_model, PCM_BYTES), + model=requested_model, + response_format=None, + ) + + assert response.response.content[:4] == b"RIFF" + assert response.response.headers["content-type"] == "audio/wav" diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 20ed8ac0c10..7f6e955e831 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -31,6 +31,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( encode_file_id_with_model, ) from litellm.types.llms.openai import OpenAIFileObject +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder class _InMemoryManagedFileTable: @@ -451,6 +452,20 @@ async def test_async_pre_call_deployment_hook_prefers_top_level_model_info(): ), "Should prefer top-level model_info over litellm_metadata" +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_reads_model_info_that_forwards_its_fields(): + proxy_managed_files: Final = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()) + kwargs: Final = { + "input_file_id": "managed-file-abc", + "model_file_id_mapping": {"managed-file-abc": {"deployment-xyz": "file-provider-123"}}, + "model_info": AttributeForwarder({"id": "deployment-xyz"}), + } + + result: Final = await proxy_managed_files.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acreate_batch) + + assert result["input_file_id"] == "file-provider-123" + + @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_no_model_info_leaves_file_id_unchanged(): """ diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index b4485bcc5f1..60600afab84 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -1921,6 +1921,62 @@ async def test_list_tools_classifies_upstream_status_for_post_failures(failure_m assert rate_limit is None +@pytest.mark.asyncio +@pytest.mark.parametrize("paginated", (False, True)) +async def test_sse_catalog_limit_reports_the_failure_and_allows_smaller_pages(paginated: bool) -> None: + from mcp.types import Tool + + limit: Final = 1024 + tools: Final = tuple( + Tool(name=f"tool_{index}", description="private-upstream-secret " * 25, input_schema={"type": "object"}) + for index in range(2) + ) + catalog: Final = tuple(tool.model_dump(by_alias=True, exclude_none=True) for tool in tools) + assert len(json.dumps({"tools": catalog}).encode()) > limit + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method != "POST": + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + assert isinstance(payload, (JSONRPCRequest, JSONRPCNotification)) + if isinstance(payload, JSONRPCNotification): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response(200, json={ + "jsonrpc": "2.0", "id": payload.id, "result": { + "protocolVersion": LATEST_HANDSHAKE_VERSION, + "capabilities": {"tools": {}}, "serverInfo": {"name": "catalog", "version": "1"}, + }, + }) + assert payload.method == "tools/list" + cursor: Final = (payload.params or {}).get("cursor") + page: Final = catalog[1:] if cursor == "next-page" else catalog[:1] if paginated else catalog + result: Final = { + "tools": page, + **({"nextCursor": "next-page"} if paginated and cursor is None else {}), + } + event: Final = f"data: {json.dumps({'jsonrpc': '2.0', 'id': payload.id, 'result': result})}\n\n" + assert (len(event.encode()) < limit) == paginated + return httpx2.Response(200, headers={"content-type": "text/event-stream"}, content=event) + + def transport_context(url: str, http_client: httpx2.AsyncClient) -> _TransportContext: + return streamable_http_client(url, http_client=http_client, max_sse_event_size=limit) + + client: Final = _MockTransportClient( + respond, server_url="https://upstream.example/mcp", transport_context=transport_context, + ) + if paginated: + assert await client.list_tools(raise_on_error=True) == list(tools) + return + with pytest.raises(MCPError) as caught: + await client.list_tools(raise_on_error=True) + fault: Final = classify_list_exception(caught.value) + assert fault == ServerListFault(tag="upstream_error", sse_event_limit_bytes=limit) + assert list_fault_http_status(fault) == 502 + assert fault.public_message is not None + assert "private-upstream-secret" not in fault.public_message + + @pytest.mark.asyncio async def test_http_status_check_allows_auth_refresh_before_rejecting() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( diff --git a/tests/unit/fine_tuning/test_main.py b/tests/unit/fine_tuning/test_main.py index 57c845c6e1a..89f51d2e74a 100644 --- a/tests/unit/fine_tuning/test_main.py +++ b/tests/unit/fine_tuning/test_main.py @@ -1,11 +1,12 @@ import json -from typing import Final +from typing import Final, Literal import httpx import pytest import respx import litellm +from litellm.types.utils import LiteLLMFineTuningJob API_BASE: Final = "https://fine-tuning.test" BASE_MODEL: Final = "gpt-4.1-nano-2025-04-14" @@ -52,3 +53,42 @@ async def test_create_sends_the_model_without_the_provider_prefix( ) assert json.loads(sent.calls.last.request.content)["model"] == BASE_MODEL + + +@pytest.mark.asyncio +@respx.mock +@pytest.mark.parametrize( + ("custom_llm_provider", "api_base", "route", "sent_api_version"), + [ + ("openai", f"{API_BASE}/v1", f"{API_BASE}/v1/fine_tuning/jobs", None), + ("azure", API_BASE, f"{API_BASE}/openai/fine_tuning/jobs", "2024-10-21"), + ], +) +@pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"]) +async def test_create_with_api_version_reaches_the_provider_and_only_azure_gets_it( + monkeypatch: pytest.MonkeyPatch, + custom_llm_provider: Literal["openai", "azure"], + api_base: str, + route: str, + sent_api_version: str | None, + is_async: bool, +) -> None: + sent: Final = respx.post(url__startswith=route).mock(return_value=httpx.Response(200, json=PROVIDER_JOB)) + request: Final = { + "model": BASE_MODEL, + "training_file": "file-abc123", + "custom_llm_provider": custom_llm_provider, + "api_key": "sk-test", + "api_base": api_base, + "api_version": "2024-10-21", + } + + async with httpx.AsyncClient() as session: + monkeypatch.setattr(litellm, "aclient_session", session) + job: Final = ( + await litellm.acreate_fine_tuning_job(**request) if is_async else litellm.create_fine_tuning_job(**request) + ) + + assert isinstance(job, LiteLLMFineTuningJob) + assert job.id == PROVIDER_JOB["id"] + assert sent.calls.last.request.url.params.get("api-version") == sent_api_version diff --git a/tests/unit/harness/test_endpoint.py b/tests/unit/harness/test_endpoint.py index 19bd6d0450e..d1447fd40cc 100644 --- a/tests/unit/harness/test_endpoint.py +++ b/tests/unit/harness/test_endpoint.py @@ -272,6 +272,7 @@ async def test_sdk_chat_stream(monkeypatch: pytest.MonkeyPatch) -> None: assert lines[-1] == "data: [DONE]" assert json.loads(lines[0][6:])["choices"][0]["delta"]["content"] == "he" assert calls[0]["stream_options"] == {"include_usage": True} + assert calls[0]["_litellm_stream_usage_injected"] is True assert (ep.usage.input_tokens, ep.usage.output_tokens) == (6, 2) assert ep.usage.cost == pytest.approx(6 * 0.001 + 2 * 0.002) diff --git a/tests/unit/integrations/dotprompt/test_prompt_manager.py b/tests/unit/integrations/dotprompt/test_prompt_manager.py index dbd4e4a4c4c..cbf19fe7985 100644 --- a/tests/unit/integrations/dotprompt/test_prompt_manager.py +++ b/tests/unit/integrations/dotprompt/test_prompt_manager.py @@ -7,6 +7,7 @@ from fastapi.testclient import TestClient +from typing import Final from unittest.mock import MagicMock, Mock, patch import httpx @@ -362,6 +363,22 @@ def test_load_prompts_from_json_data(): assert rendered == "Dynamic prompt: test" +def test_load_prompts_from_json_data_skips_malformed_entries(): + manager: Final = PromptManager(prompt_directory=str(Path(__file__).parent)) + initial_count: Final = len(manager.prompts) + + manager.load_prompts_from_json_data( + { + "valid_prompt": {"content": "Valid: {{value}}", "metadata": {"model": "gpt-4"}}, + "string_entry": "not a prompt", + "null_entry": None, + } + ) + + assert len(manager.prompts) == initial_count + 1 + assert manager.render("valid_prompt", {"value": "ok"}) == "Valid: ok" + + def test_prompt_file_to_json_conversion(): """Test converting .prompt files to JSON format.""" # Create a temporary prompt file with frontmatter @@ -661,6 +678,27 @@ def test_prompt_initializer_registers_flat_db_prompt_under_base_id(): assert template.content == "AHOY {{name}}" +def test_prompt_initializer_accepts_params_without_dotprompt_content(): + from types import SimpleNamespace + + from litellm.integrations.dotprompt import prompt_initializer + from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec + + litellm_params: Final = SimpleNamespace( + prompt_data={"content": "AHOY {{name}}", "metadata": {"model": "gpt-4o-mini"}}, + ) + prompt_spec: Final = PromptSpec( + prompt_id="agent-prompt", + litellm_params=PromptLiteLLMParams(prompt_integration="dotprompt"), + ) + + dotprompt_manager: Final = prompt_initializer(litellm_params, prompt_spec) + + template: Final = dotprompt_manager.prompt_manager.get_prompt("agent-prompt") + assert template is not None + assert template.content == "AHOY {{name}}" + + def _swap_prompt_manager_and_spec(ignore_prompt_manager_model: bool) -> tuple[DotpromptManager, PromptSpec]: manager = DotpromptManager( prompt_data={"content": "You are a pirate assistant.", "metadata": {"model": "gpt-4o-mini"}}, diff --git a/tests/unit/integrations/newrelic/test_newrelic.py b/tests/unit/integrations/newrelic/test_newrelic.py index 00a115dd8f2..d09732c6eb2 100644 --- a/tests/unit/integrations/newrelic/test_newrelic.py +++ b/tests/unit/integrations/newrelic/test_newrelic.py @@ -1,6 +1,7 @@ import os import sys from datetime import datetime, timezone +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -491,6 +492,18 @@ class TestExtractAllMessagesTimestamps: input_msg = next(m for m in messages if not m.get("is_response")) assert input_msg["timestamp"] == int(1_000_000.0 * 1000.0) + def test_input_messages_read_a_start_time_that_only_exposes_timestamp(self): + class EpochSeconds: + def timestamp(self) -> float: + return 1_000_000.0 + + kwargs: Final = {**make_kwargs(messages=[{"role": "user", "content": "Hi"}]), "start_time": EpochSeconds()} + + messages: Final = self.logger._extract_all_messages(kwargs, make_response(), response_model="gpt-4", vendor="openai") + + input_msg: Final = next(m for m in messages if not m.get("is_response")) + assert input_msg["timestamp"] == int(1_000_000.0 * 1000.0) + def test_output_messages_get_end_time_timestamp(self): kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}]) response = make_response() diff --git a/tests/unit/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py index 3f4389765d3..8af13fd6e7e 100644 --- a/tests/unit/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -14,6 +14,7 @@ from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer import pytest import requests +from requests_mock import Adapter pytest.importorskip("opentelemetry") @@ -22,16 +23,18 @@ from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( # noqa: ) from opentelemetry import baggage # noqa: E402 from opentelemetry.context import attach, detach # noqa: E402 -from opentelemetry._logs.severity import SeverityNumber # noqa: E402 -from opentelemetry.sdk._logs import LogData, LogRecord # noqa: E402 -from opentelemetry.sdk._logs.export import LogExportResult # noqa: E402 +from opentelemetry._logs import LogRecord, SeverityNumber # noqa: E402 # optional SDK checked above +from opentelemetry.sdk._logs import ReadableLogRecord # noqa: E402 # optional SDK checked above +from opentelemetry.sdk._logs.export import LogRecordExportResult # noqa: E402 # optional SDK checked above +from opentelemetry.sdk.resources import Resource # noqa: E402 # optional SDK checked above from opentelemetry.sdk.metrics import MeterProvider # noqa: E402 from opentelemetry.sdk.metrics.export import InMemoryMetricReader # noqa: E402 -from opentelemetry.sdk.trace import TracerProvider # noqa: E402 +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider # noqa: E402 # optional SDK checked above from opentelemetry.sdk.trace.export import ( # noqa: E402 BatchSpanProcessor, ConsoleSpanExporter, SimpleSpanProcessor, + SpanExportResult, ) from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402 InMemorySpanExporter, @@ -46,6 +49,7 @@ import litellm # noqa: E402 from tests.unit.integrations.conftest import TlsSink # noqa: E402 from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 +from litellm.integrations.otel.plumbing.otlp_json import OTLPJsonSpanExporter # noqa: E402 # optional SDK checked above from litellm.integrations.otel.model.config import OpenTelemetryV2Config # noqa: E402 from litellm.integrations.otel.mappers.genai import GenAIMapper # noqa: E402 from litellm.integrations.otel.mappers.legacy import LegacyMapper # noqa: E402 @@ -733,7 +737,9 @@ def _export_one_trace_to_local_collector(exporter_kind: str) -> tuple[list[dict] class Collector(BaseHTTPRequestHandler): def do_POST(self): body = self.rfile.read(int(self.headers["Content-Length"])) - received.append({"path": self.path, "headers": dict(self.headers), "body": body}) + received.append( + {"path": self.path, "headers": requests.structures.CaseInsensitiveDict(self.headers), "body": body} + ) self.send_response(200) self.end_headers() @@ -793,6 +799,95 @@ def test_http_json_exporter_posts_otlp_json_to_traces_endpoint(): assert child["parentSpanId"] == format(parent_id, "016x") +@pytest.mark.parametrize( + ("statuses", "expected"), + [ + ((200,), SpanExportResult.SUCCESS), + ((400,), SpanExportResult.FAILURE), + ((503, 200), SpanExportResult.SUCCESS), + ], +) +def test_http_json_exporter_preserves_http_failure_and_retry_results( + statuses: tuple[int, ...], expected: SpanExportResult +) -> None: + adapter: Final = Adapter() + endpoint: Final = "https://collector.test/v1/traces" + adapter.register_uri( + "POST", + endpoint, + [{"status_code": status, "headers": {"Retry-After": "0"}} for status in statuses], + ) + with requests.Session() as session: + session.mount("https://", adapter) + exporter: Final = OTLPJsonSpanExporter( + endpoint=endpoint, + headers={ + "CONTENT-TYPE": "application/x-protobuf", + "x-collector-token": "secret", + }, + session=session, + ) + try: + assert exporter.export(()) is expected + finally: + exporter.shutdown() + assert len(adapter.request_history) == len(statuses) + for request in adapter.request_history: + assert request.headers["Content-Type"] == "application/json" + assert request.headers["x-collector-token"] == "secret" + assert request.json() == {"resourceSpans": []} + + +@pytest.mark.parametrize( + ("size_limit", "expected", "request_count"), + [ + (0, SpanExportResult.SUCCESS, 1), + (len(b'{"resourceSpans":[]}'), SpanExportResult.SUCCESS, 1), + (len(b'{"resourceSpans":[]}') - 1, SpanExportResult.FAILURE, 0), + ], +) +def test_http_json_exporter_checks_json_size_before_sending( + size_limit: int, expected: SpanExportResult, request_count: int +) -> None: + adapter: Final = Adapter() + endpoint: Final = "https://collector.test/v1/traces" + adapter.register_uri("POST", endpoint, status_code=200) + with requests.Session() as session: + session.mount("https://", adapter) + exporter: Final = OTLPJsonSpanExporter(endpoint, {}, session=session, max_request_size=size_limit) + try: + assert exporter.export(()) is expected + finally: + exporter.shutdown() + assert len(adapter.request_history) == request_count + + +def test_http_json_exporter_rejects_export_after_shutdown() -> None: + adapter: Final = Adapter() + endpoint: Final = "https://collector.test/v1/traces" + adapter.register_uri("POST", endpoint, status_code=200) + with requests.Session() as session: + session.mount("https://", adapter) + exporter: Final = OTLPJsonSpanExporter(endpoint, {}, session=session) + exporter.shutdown() + assert exporter.export(()) is SpanExportResult.FAILURE + assert adapter.request_history == [] + + +def test_http_json_exporter_returns_failure_for_unencodable_spans() -> None: + adapter: Final = Adapter() + endpoint: Final = "https://collector.test/v1/traces" + adapter.register_uri("POST", endpoint, status_code=200) + with requests.Session() as session: + session.mount("https://", adapter) + exporter: Final = OTLPJsonSpanExporter(endpoint, {}, session=session) + try: + assert exporter.export((ReadableSpan("missing-context"),)) is SpanExportResult.FAILURE + finally: + exporter.shutdown() + assert adapter.request_history == [] + + def test_http_protobuf_exporter_still_posts_protobuf(): (_child_request, parent_request), (trace_id, _parent_id, _child_id) = _export_one_trace_to_local_collector( "http/protobuf" @@ -1630,9 +1725,13 @@ def test_v2_otlp_http_log_export_trusts_ssl_cert_file(monkeypatch: pytest.Monkey severity_number=SeverityNumber.INFO, body="v2-tls-test", ) - log_data = LogData(log_record=record, instrumentation_scope=InstrumentationScope("v2-tls-test")) + log_data: Final = ReadableLogRecord( + log_record=record, + resource=Resource({}), + instrumentation_scope=InstrumentationScope("v2-tls-test"), + ) result = exporter.export([log_data]) - assert result is LogExportResult.SUCCESS, f"log export failed: {result}" + assert result is LogRecordExportResult.SUCCESS, f"log export failed: {result}" assert tls_sink.received.get(timeout=5) == "/v1/logs" finally: exporter.shutdown() diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index c7a057cca5c..782442aad22 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -4,6 +4,7 @@ import asyncio import contextvars import threading import time +import uuid from base64 import b64encode from collections.abc import Mapping from dataclasses import replace @@ -14,10 +15,12 @@ from typing import Final import pytest from opentelemetry import trace as trace_api +from opentelemetry.context import Context from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.id_generator import RandomIdGenerator from opentelemetry.trace import Status, StatusCode import litellm @@ -1159,12 +1162,36 @@ class TestFanOut: ARIZE_TEAM_PARAMS = {"arize_space_id": "space-team", "arize_api_key": "key-team"} +def caller_context(trace_id: int) -> Context: + return trace_api.set_span_in_context( + trace_api.NonRecordingSpan( + trace_api.SpanContext(trace_id=trace_id, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01)) + ) + ) + + def arize_destination(**rates: str) -> OtelDestination: destination = destination_for("arize", {**ARIZE_TEAM_PARAMS, **rates}) assert destination is not None, "the fixture must resolve for the test to mean anything" return destination +#: Trace ids whose SHA-256 reads 0.40 and 0.80 of the way through its first eight bytes, +#: so a rate of 0.5 keeps the first and drops the second. +KEPT_AT_HALF: Final = 4 +DROPPED_AT_HALF: Final = 5 + + +class _FixedTraceId(RandomIdGenerator): + """Every trace the provider starts gets ``trace_id``.""" + + def __init__(self, trace_id: int) -> None: + self._trace_id = trace_id + + def generate_trace_id(self) -> int: + return self._trace_id + + class TestTenantSampling: """A team's ``arize_success_sampling_rate`` and ``arize_error_sampling_rate`` decide how much of its traffic the fan-out delivers to its space, the way they already decide what @@ -1173,14 +1200,17 @@ class TestTenantSampling: NAMES = {"POST /v1/chat/completions", "auth /v1/chat/completions", "chat gpt-4"} @staticmethod - def _fan_out(dest_exporter, draw=None, global_exporter=None): - provider = TracerProvider() + def _fan_out( + dest_exporter: InMemorySpanExporter, + trace_id: int | None = None, + global_exporter: InMemorySpanExporter | None = None, + ) -> TracerProvider: + provider: Final = TracerProvider() if trace_id is None else TracerProvider(id_generator=_FixedTraceId(trace_id)) if global_exporter is not None: provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(global_exporter), "arize")) - kwargs = {"processor_factory": lambda _d: SimpleSpanProcessor(dest_exporter)} - if draw is not None: - kwargs["sampling_draw"] = draw - provider.add_span_processor(TenantFanOutSpanProcessor(**kwargs)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) return provider @staticmethod @@ -1224,7 +1254,7 @@ class TestTenantSampling: def test_an_unset_rate_exports_everything(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.99) + provider: Final = self._fan_out(dest_exporter, trace_id=DROPPED_AT_HALF) self._run(provider, (arize_destination(),)) @@ -1300,10 +1330,10 @@ class TestTenantSampling: def test_a_span_that_ends_after_the_root_follows_the_requests_verdict(self): """A post-call database write that starts after the whole request tree has ended - goes where the rest of the tree went, without a second draw.""" + goes where the rest of the tree went, without a second verdict.""" kept_exporter, dropped_exporter = InMemorySpanExporter(), InMemorySpanExporter() - kept = self._fan_out(kept_exporter, draw=lambda: 0.1) - dropped = self._fan_out(dropped_exporter, draw=lambda: 0.9) + kept: Final = self._fan_out(kept_exporter, trace_id=KEPT_AT_HALF) + dropped: Final = self._fan_out(dropped_exporter, trace_id=DROPPED_AT_HALF) destinations = (arize_destination(arize_success_sampling_rate="0.5"),) def late_tree(provider): @@ -1327,14 +1357,263 @@ class TestTenantSampling: } assert dropped_exporter.get_finished_spans() == () + def test_a_span_that_ends_after_a_failed_root_follows_the_failed_requests_verdict(self): + """A post-call database write after a failed request goes where that request went, + though on its own it would answer to the success rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + root_context: Final = trace_api.set_span_in_context(root) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=root_context): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "postgres INSERT LiteLLM_SpendLogs", + } + + def test_a_failed_call_after_a_dropped_call_of_its_trace_still_answers_to_the_error_rate(self): + """A process that dropped a successful call of the trace keeps a later failed call of + it at an unset error rate, as a process that only served the failed call does.""" + served_both, served_failed = InMemorySpanExporter(), InMemorySpanExporter() + destinations: Final = (arize_destination(arize_success_sampling_rate="0.5"),) + both: Final = self._fan_out(served_both, trace_id=DROPPED_AT_HALF) + + self._run(both, destinations) + self._run(both, destinations, failed=True) + self._run(self._fan_out(served_failed, trace_id=DROPPED_AT_HALF), destinations, failed=True) + + assert [sorted(s.name for s in exporter.get_finished_spans()) for exporter in (served_both, served_failed)] == [ + sorted(self.NAMES) + ] * 2 + + def test_a_successful_call_after_a_kept_failed_call_of_its_trace_answers_to_the_success_rate(self): + """A process that kept a failed call of the trace drops a later successful call of it + at a zero success rate, as a process that only served the successful call does.""" + served_both, served_success = InMemorySpanExporter(), InMemorySpanExporter() + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + both: Final = self._fan_out(served_both, trace_id=KEPT_AT_HALF) + + self._run(both, destinations, failed=True) + self._run(both, destinations) + self._run(self._fan_out(served_success, trace_id=KEPT_AT_HALF), destinations) + + assert len(served_both.get_finished_spans()) == len(self.NAMES), "only the failed call" + assert all(s.name in self.NAMES for s in served_both.get_finished_spans()) + assert served_success.get_finished_spans() == () + + def test_a_late_child_that_fails_after_its_request_was_dropped_is_dropped_with_it(self): + """A post-call write that fails after its successful request was dropped goes where + the request went, rather than reaching the space on its own at the error rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + root_context: Final = trace_api.set_span_in_context(root) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=root_context) as write: + write.set_status(Status(StatusCode.ERROR, "connection reset")) + + in_fresh_context(run) + + assert dest_exporter.get_finished_spans() == () + + def test_a_failure_deep_in_a_request_marks_the_whole_request(self): + """A database span that fails under the auth span puts the whole request, root + included, on the error rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("auth /v1/chat/completions"): + with tracer.start_as_current_span("postgres SELECT LiteLLM_VerificationToken") as lookup: + lookup.set_status(Status(StatusCode.ERROR, "connection reset")) + with tracer.start_as_current_span("chat gpt-4"): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + *self.NAMES, + "postgres SELECT LiteLLM_VerificationToken", + } + + def test_a_span_two_levels_down_that_outlives_many_other_spans_still_closes_its_request(self): + """A classifier call under the route span that stays open while the worker serves + more requests than it tracks at once still counts against its own request, so the + request is decided once it ends, on its own outcome.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def other_request(): + set_request_destinations(()) + with get_tracer(provider, "litellm").start_as_current_span("POST /v1/chat/completions"): + pass + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + with tracer.start_as_current_span("route gpt-4"): + with tracer.start_as_current_span("chat classifier"): + for _ in range(otel_providers._MAX_OPEN_REQUESTS + 1): + contextvars.Context().run(other_request) + with tracer.start_as_current_span("chat gpt-4"): + pass + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "chat classifier", + "chat gpt-4", + } + + def test_a_failure_two_levels_down_reaches_its_request_however_many_spans_are_open(self): + """A classifier call that fails while another request holds more spans open than + the worker tracks requests marks its own request, which is then kept at the error + rate as a whole.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + others: Final = contextvars.Context() + + def open_other_request(): + set_request_destinations(()) + tracer: Final = get_tracer(provider, "litellm") + root: Final = tracer.start_span("POST /v1/chat/completions") + parent: Final = trace_api.set_span_in_context(root) + return ( + root, + *(tracer.start_span("chat flood", context=parent) for _ in range(otel_providers._MAX_OPEN_REQUESTS)), + ) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("route gpt-4"): + with tracer.start_as_current_span("chat classifier") as classifier: + held_open: Final = others.run(open_other_request) + classifier.set_status(Status(StatusCode.ERROR, "classifier timed out")) + return held_open + + held_open: Final = in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "chat classifier", + } + for span in held_open: + span.end() + + def test_the_fan_out_forgets_a_span_once_it_ends_and_a_request_once_it_is_pushed_out(self): + """The request lookup holds only spans still open, and only those of the requests + the fan-out still counts, so neither ended spans nor a flood of open requests grow + it past the open request bound.""" + fan_out: Final = TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(InMemorySpanExporter()) + ) + provider: Final = TracerProvider() + provider.add_span_processor(fan_out) + + def run(): + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("chat gpt-4"): + pass + assert (fan_out._roots, dict(fan_out._open)) == ({}, {}) + root: Final = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) + tracer.start_span("route gpt-4", context=trace_api.set_span_in_context(root)) + return tuple( + tracer.start_span(f"POST {index}", context=trace_api.Context()) + for index in range(otel_providers._MAX_OPEN_REQUESTS) + ) + + held_open: Final = in_fresh_context(run) + + assert len(fan_out._roots) == otel_providers._MAX_OPEN_REQUESTS + for span in held_open: + span.end() + assert (fan_out._roots, dict(fan_out._open)) == ({}, {}) + + def test_a_span_starting_under_an_ended_span_below_the_root_follows_its_request(self): + """A post-call write started under the route span after its failed request ended + is kept with the request at the error rate, rather than decided on its own at the + success rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + with tracer.start_as_current_span("route gpt-4") as route: + route_context: Final = trace_api.set_span_in_context(route) + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=route_context): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "postgres INSERT LiteLLM_SpendLogs", + } + + def test_concurrent_calls_of_one_trace_each_answer_to_their_own_outcome(self): + """Two calls of the caller's trace in flight at once on one process, one failing, + are decided as two processes serving one each would decide them.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + remote: Final = trace_api.SpanContext( + trace_id=KEPT_AT_HALF, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01) + ) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + caller: Final = trace_api.set_span_in_context(trace_api.NonRecordingSpan(remote)) + with tracer.start_as_current_span("POST failed", context=caller) as failed: + with tracer.start_as_current_span("POST succeeded", context=caller): + pass + failed.set_status(Status(StatusCode.ERROR, "upstream 500")) + + in_fresh_context(run) + + assert [s.name for s in dest_exporter.get_finished_spans()] == ["POST failed"] + def test_a_late_span_whose_verdict_was_forgotten_is_decided_on_its_own_not_stranded(self): """Once enough other requests have been decided to evict a request's verdict, a span of it that starts late (a post-call database write) opens and closes its own count, - so it is decided on a draw of its own as soon as it ends instead of waiting in a tree - nothing would ever close.""" + so it is decided on its own as soon as it ends instead of waiting in a tree nothing + would ever close.""" dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.1) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1353,18 +1632,17 @@ class TestTenantSampling: assert names[0] == "POST /v1/chat/completions" assert names[-1] == "postgres INSERT LiteLLM_SpendLogs", "exported as it ended, not held for a root" - def test_a_callers_traceparent_neither_steers_the_draw_nor_hides_the_root(self): - """The draw is random, not read from the trace id a caller may have chosen, and the - server span under a remote parent still closes the tree.""" + def test_a_callers_traceparent_decides_by_its_trace_id_and_still_closes_the_tree(self): + """A request continuing a caller's trace is decided on that trace id, and the server + span under a remote parent still closes the tree.""" dest_exporter = InMemorySpanExporter() - draws = iter([0.9, 0.1]) - provider = self._fan_out(dest_exporter, draw=lambda: next(draws)) + provider: Final = self._fan_out(dest_exporter) destinations = (arize_destination(arize_success_sampling_rate="0.5"),) def run(): set_request_destinations(destinations) tracer = get_tracer(provider, "litellm") - for trace_id in (1, 2): + for trace_id in (DROPPED_AT_HALF, KEPT_AT_HALF): remote = trace_api.SpanContext( trace_id=trace_id, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01) ) @@ -1377,12 +1655,12 @@ class TestTenantSampling: kept = dest_exporter.get_finished_spans() assert [s.name for s in kept] == ["chat gpt-4", "POST /v1/chat/completions"] - assert all(s.context.trace_id == 2 for s in kept), "the second request, whose draw of 0.1 passed" + assert all(s.context.trace_id == KEPT_AT_HALF for s in kept) def test_a_root_that_never_ends_holds_the_tree_only_up_to_the_bound(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1397,8 +1675,8 @@ class TestTenantSampling: def test_a_flood_of_waiting_trees_decides_the_oldest_on_what_it_has(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1413,21 +1691,21 @@ class TestTenantSampling: assert [s.name for s in dest_exporter.get_finished_spans()] == ["chat 0"] - def test_a_flood_of_open_traces_forgets_the_oldest_and_decides_it_as_its_spans_end(self): - """The fan-out counts open spans for a bounded number of traces, a bound far above - what one process keeps in flight, since every trace of the provider is counted, not - only the sampled ones. A trace pushed out of that count is decided on what it has as + def test_a_flood_of_open_requests_forgets_the_oldest_and_decides_it_as_its_spans_end(self): + """The fan-out counts open spans for a bounded number of requests, a bound far above + what one process keeps in flight, since every request of the provider is counted, not + only the sampled ones. A request pushed out of that count is decided on what it has as its next span ends, root still open, rather than held until a bound gets to it.""" dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) tracer = get_tracer(provider, "litellm") root = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) tracer.start_span("auth /v1/chat/completions", context=trace_api.set_span_in_context(root)).end() - for index in range(otel_providers._MAX_OPEN_TRACES): + for index in range(otel_providers._MAX_OPEN_REQUESTS): tracer.start_span(f"POST {index}", context=trace_api.Context()) assert dest_exporter.get_finished_spans() == (), "still held before the count is forgotten" tracer.start_span("chat gpt-4", context=trace_api.set_span_in_context(root)).end() @@ -1443,14 +1721,12 @@ class TestTenantSampling: def test_shutdown_decides_what_is_still_held(self): dest_exporter = InMemorySpanExporter() - fan_out = TenantFanOutSpanProcessor( - processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter), sampling_draw=lambda: 0.0 - ) + fan_out: Final = TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) provider = TracerProvider() provider.add_span_processor(fan_out) def run(): - set_request_destinations((arize_destination(arize_success_sampling_rate="0.5"),)) + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) tracer = get_tracer(provider, "litellm") with tracer.start_as_current_span("POST /v1/chat/completions"): with tracer.start_as_current_span("chat gpt-4"): @@ -1476,14 +1752,12 @@ class TestTenantSampling: straggler.end() super().on_end(span) - fan_out = TenantFanOutSpanProcessor( - processor_factory=lambda _d: _EndsAStragglerOnExport(dest_exporter), sampling_draw=lambda: 0.0 - ) + fan_out: Final = TenantFanOutSpanProcessor(processor_factory=lambda _d: _EndsAStragglerOnExport(dest_exporter)) provider = TracerProvider() provider.add_span_processor(fan_out) def run(): - set_request_destinations((arize_destination(arize_success_sampling_rate="0.5"),)) + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) tracer = get_tracer(provider, "litellm") root = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) tracer.start_span("chat gpt-4", context=trace_api.set_span_in_context(root)).end() @@ -1495,20 +1769,20 @@ class TestTenantSampling: assert {s.name for s in dest_exporter.get_finished_spans()} == {"chat gpt-4", "embed ada"} - def test_a_zero_rate_drops_even_a_draw_of_zero(self): + def test_a_zero_rate_drops_a_trace_a_half_rate_keeps(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) + provider: Final = self._fan_out(dest_exporter, trace_id=KEPT_AT_HALF) self._run(provider, (arize_destination(arize_success_sampling_rate="0.0"),)) assert dest_exporter.get_finished_spans() == () - def test_the_draw_is_compared_to_the_rate_the_way_the_legacy_callback_compares_it(self): + def test_the_trace_id_is_read_against_the_rate(self): dropped_exporter, kept_exporter = InMemorySpanExporter(), InMemorySpanExporter() - destinations = (arize_destination(arize_success_sampling_rate="0.2"),) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.5"),) - self._run(self._fan_out(dropped_exporter, draw=lambda: 0.3), destinations) - self._run(self._fan_out(kept_exporter, draw=lambda: 0.2), destinations) + self._run(self._fan_out(dropped_exporter, trace_id=DROPPED_AT_HALF), destinations) + self._run(self._fan_out(kept_exporter, trace_id=KEPT_AT_HALF), destinations) assert dropped_exporter.get_finished_spans() == () assert {s.name for s in kept_exporter.get_finished_spans()} == self.NAMES @@ -1527,6 +1801,39 @@ class TestTenantSampling: assert all(names == self.NAMES for names in by_trace.values()) assert 0 < len(by_trace) < 64 + @pytest.mark.parametrize( + ("rates", "failed"), + [({"arize_success_sampling_rate": "0.5"}, False), ({"arize_error_sampling_rate": "0.5"}, True)], + ids=["success-rate", "error-rate"], + ) + def test_two_processes_serving_one_trace_reach_the_same_verdict(self, rates: Mapping[str, str], failed: bool): + """Two pods behind a load balancer each serve one call of a trace the caller + continued with its ``traceparent``: the team's space gets both calls or neither, + and a rate of 0.5 over UUID-shaped trace ids keeps some traces and drops others.""" + exporters: Final = (InMemorySpanExporter(), InMemorySpanExporter()) + pods: Final = tuple(self._fan_out(exporter) for exporter in exporters) + destinations: Final = (arize_destination(**rates),) + trace_ids: Final = tuple(uuid.uuid5(uuid.NAMESPACE_URL, f"pod/{index}").int for index in range(32)) + + def run(provider: TracerProvider): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + for trace_id in trace_ids: + with tracer.start_as_current_span("POST /v1/chat/completions", context=caller_context(trace_id)): + with tracer.start_as_current_span("chat gpt-4") as span: + if failed: + span.set_status(Status(StatusCode.ERROR, "upstream 500")) + + for provider in pods: + in_fresh_context(run, provider) + + first, second = ( + {s.context.trace_id for s in exporter.get_finished_spans() if s.context is not None} + for exporter in exporters + ) + assert first == second + assert 0 < len(first) < len(trace_ids) + def test_in_additive_mode_the_operators_own_copy_is_not_sampled(self, monkeypatch): """The rates are the team's setting for the team's space; the operator's backbone keeps every span.""" @@ -1547,7 +1854,7 @@ class TestTenantSampling: assert {s.name for s in dest_exporter.get_finished_spans()} == self.NAMES - def test_two_views_of_one_exporter_each_draw_at_their_own_rate(self): + def test_two_views_of_one_exporter_each_answer_to_their_own_rate(self): """Two destinations to the same space with different rates share an exporter, not a verdict: the 1.0 view gets the tree once and the 0.0 view adds nothing.""" dest_exporter = InMemorySpanExporter() @@ -1582,7 +1889,7 @@ class TestTenantSampling: @pytest.mark.parametrize("bad", ["abc", "-0.1", "1.5", "nan"]) def test_an_unusable_rate_exports_rather_than_dropping(self, bad): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.99) + provider: Final = self._fan_out(dest_exporter, trace_id=DROPPED_AT_HALF) self._run(provider, (arize_destination(arize_success_sampling_rate=bad),)) diff --git a/tests/unit/integrations/otel/test_otel_v2_mount.py b/tests/unit/integrations/otel/test_otel_v2_mount.py index 3a27c3ec4e4..61f8b9ae0d9 100644 --- a/tests/unit/integrations/otel/test_otel_v2_mount.py +++ b/tests/unit/integrations/otel/test_otel_v2_mount.py @@ -382,7 +382,9 @@ def test_every_traced_response_names_its_own_server_span(monkeypatch): assert responses["/stream"].text == "data: first\n\ndata: [DONE]\n\n" for path, response in responses.items(): context = server_spans[path].context - assert response.headers.get_list("traceresponse") == [f"00-{context.trace_id:032x}-{context.span_id:016x}-01"] + assert response.headers.get_list("traceresponse") == [ + f"00-{context.trace_id:032x}-{context.span_id:016x}-{context.trace_flags:02x}" + ] @pytest.mark.parametrize("path", ["/relayed", "/relayed/health"]) @@ -545,4 +547,4 @@ def test_the_unhandled_error_helper_names_the_span_it_is_given_in_any_context(mo headers = contextvars.Context().run(trace_response_headers, span) - assert headers["traceresponse"] == f"00-{context.trace_id:032x}-{context.span_id:016x}-01" + assert headers["traceresponse"] == f"00-{context.trace_id:032x}-{context.span_id:016x}-{context.trace_flags:02x}" diff --git a/tests/unit/integrations/otel/test_otel_v2_presets.py b/tests/unit/integrations/otel/test_otel_v2_presets.py index 64366626784..570c27a24e1 100644 --- a/tests/unit/integrations/otel/test_otel_v2_presets.py +++ b/tests/unit/integrations/otel/test_otel_v2_presets.py @@ -2,13 +2,19 @@ event loop: the preset does no network I/O, and a custom exporter mints the JWT lazily on its first export (in the BatchSpanProcessor worker thread).""" +import json + import httpx import pytest +import requests_mock +import respx +from opentelemetry.sdk.trace.export import SpanExportResult from litellm.integrations.otel.plumbing import providers from litellm.integrations.otel.model.config import ExporterSpec from litellm.integrations.otel.presets import agentops as agentops_mod from litellm.integrations.otel.presets.agentops import ( + _AGENTOPS_AUTH_ENDPOINT, _AGENTOPS_ENDPOINT, _AGENTOPS_EXPORTER_KIND, _build_agentops_exporter, @@ -75,50 +81,44 @@ def test_dynamic_cred_presets_tag_exporter_with_matching_owner(monkeypatch): ) -def test_agentops_exporter_mints_jwt_lazily(monkeypatch): - pytest.importorskip("opentelemetry.exporter.otlp.proto.http.trace_exporter") - monkeypatch.setattr( - agentops_mod, "_fetch_agentops_jwt", lambda _k: {"token": "jwt-xyz"} - ) - spec = ExporterSpec( - kind=_AGENTOPS_EXPORTER_KIND, - endpoint=_AGENTOPS_ENDPOINT, - options={"api_key": "ak"}, - ) - exporter = _build_agentops_exporter(spec) - - # No auth header until the first export triggers the (off-loop) fetch. - assert "Authorization" not in exporter._session.headers - exporter._ensure_authenticated() - assert exporter._session.headers["Authorization"] == "Bearer jwt-xyz" - - # Cached: a second resolution does not re-fetch. - calls = [] - monkeypatch.setattr( - agentops_mod, - "_fetch_agentops_jwt", - lambda k: calls.append(k) or {"token": "again"}, - ) - exporter._ensure_authenticated() - assert calls == [] - - -def test_agentops_exporter_tolerates_fetch_failure(monkeypatch): - pytest.importorskip("opentelemetry.exporter.otlp.proto.http.trace_exporter") - - def _raise(_k): - raise RuntimeError("auth down") - - monkeypatch.setattr(agentops_mod, "_fetch_agentops_jwt", _raise) - exporter = _build_agentops_exporter( +def _agentops_exporter(): + return _build_agentops_exporter( ExporterSpec( kind=_AGENTOPS_EXPORTER_KIND, endpoint=_AGENTOPS_ENDPOINT, options={"api_key": "ak"}, ) ) - exporter._ensure_authenticated() # must not raise - assert "Authorization" not in exporter._session.headers + + +def test_agentops_exporter_mints_the_jwt_once_on_first_export_and_sends_it_with_every_batch(): + pytest.importorskip("opentelemetry.exporter.otlp.proto.http.trace_exporter") + with respx.mock() as auth, requests_mock.Mocker() as collector: + token_route = auth.post(_AGENTOPS_AUTH_ENDPOINT).respond(200, json={"token": "jwt-xyz"}) + collector.post(_AGENTOPS_ENDPOINT, status_code=200) + exporter = _agentops_exporter() + assert token_route.call_count == 0, "building the exporter must not fetch the JWT" + + results = [exporter.export(()), exporter.export(())] + + assert results == [SpanExportResult.SUCCESS, SpanExportResult.SUCCESS] + assert [json.loads(call.request.content) for call in token_route.calls] == [{"api_key": "ak"}] + assert [request.headers.get("Authorization") for request in collector.request_history] == [ + "Bearer jwt-xyz", + "Bearer jwt-xyz", + ] + + +def test_agentops_exporter_still_exports_without_auth_when_the_jwt_fetch_fails(): + pytest.importorskip("opentelemetry.exporter.otlp.proto.http.trace_exporter") + with respx.mock() as auth, requests_mock.Mocker() as collector: + auth.post(_AGENTOPS_AUTH_ENDPOINT).mock(side_effect=httpx.ConnectError("auth down")) + collector.post(_AGENTOPS_ENDPOINT, status_code=200) + + result = _agentops_exporter().export(()) + + assert result is SpanExportResult.SUCCESS + assert [request.headers.get("Authorization") for request in collector.request_history] == [None] def test_fetch_jwt_uses_owned_client_not_shared_pool(monkeypatch): diff --git a/tests/unit/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py index 9f536ab0dac..c7608159515 100644 --- a/tests/unit/integrations/test_langfuse_otel.py +++ b/tests/unit/integrations/test_langfuse_otel.py @@ -513,6 +513,22 @@ class TestLangfuseOtelIntegration: # Endpoint assertion removed as side effect is gone +def _exported_auth_headers(exporter): + import requests_mock + from opentelemetry.sdk.trace.export import SpanExportResult + + with requests_mock.Mocker() as collector: + collector.post(requests_mock.ANY, status_code=200) + assert exporter.export(()) is SpanExportResult.SUCCESS + return [ + ( + request.url, + {name: request.headers[name] for name in ("Authorization", "x-langfuse-ingestion-version")}, + ) + for request in collector.request_history + ] + + class TestLangfuseOtelKeyDynamicConfig: """Key/team-scoped Langfuse credentials must define the full export target (OTLP endpoint + auth), not just auth headers on the init-time exporter.""" @@ -617,15 +633,15 @@ class TestLangfuseOtelKeyDynamicConfig: exporter = span_processors[0].span_exporter assert isinstance(exporter, OTLPSpanExporter) - assert exporter._endpoint == "https://langfuse.example.com/api/public/otel/v1/traces" - import base64 expected_auth = base64.b64encode(b"key_public:key_secret").decode() - assert exporter._headers == { - "Authorization": f"Basic {expected_auth}", - "x-langfuse-ingestion-version": "4", - } + assert _exported_auth_headers(exporter) == [ + ( + "https://langfuse.example.com/api/public/otel/v1/traces", + {"Authorization": f"Basic {expected_auth}", "x-langfuse-ingestion-version": "4"}, + ) + ] def test_key_dynamic_params_reuse_cached_provider(self): with self._clean_env(): @@ -678,10 +694,12 @@ class TestLangfuseOtelKeyDynamicConfig: provider = next(iter(logger._tracer_provider_cache.values())).provider exporter = provider._active_span_processor._span_processors[0].span_exporter assert isinstance(exporter, OTLPSpanExporter) - assert exporter._headers == { - "Authorization": f"Basic {secret}", - "x-langfuse-ingestion-version": "4", - } + assert _exported_auth_headers(exporter) == [ + ( + "https://langfuse.example.com/api/public/otel/v1/traces", + {"Authorization": f"Basic {secret}", "x-langfuse-ingestion-version": "4"}, + ) + ] class TestLangfuseOtelResponsesAPI: diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index d14703e48a1..f908c685382 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -19,14 +19,16 @@ import pytest # Adds the grandparent directory to sys.path to allow importing project modules from opentelemetry import trace +from opentelemetry._logs import LogRecord from opentelemetry._logs.severity import SeverityNumber -from opentelemetry.sdk._logs import LogData, LogRecord +from opentelemetry.sdk._logs import ReadableLogRecord from opentelemetry.sdk._logs import LoggerProvider as OTLoggerProvider -from opentelemetry.sdk._logs.export import InMemoryLogExporter, LogExportResult, SimpleLogRecordProcessor +from opentelemetry.sdk._logs.export import InMemoryLogExporter, LogRecordExportResult, SimpleLogRecordProcessor from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import InMemoryMetricReader, MetricsData +from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor, SpanExportResult from opentelemetry.sdk.util.instrumentation import InstrumentationScope from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from parameterized import parameterized @@ -2321,9 +2323,13 @@ def test_otlp_http_log_export_trusts_ssl_cert_file(monkeypatch: pytest.MonkeyPat severity_number=SeverityNumber.INFO, body="tls-export-test", ) - log_data: Final = LogData(log_record=record, instrumentation_scope=InstrumentationScope("tls-export-test")) + log_data: Final = ReadableLogRecord( + log_record=record, + resource=Resource({}), + instrumentation_scope=InstrumentationScope("tls-export-test"), + ) result: Final = exporter.export([log_data]) - assert result is LogExportResult.SUCCESS, f"log export failed: {result}" + assert result is LogRecordExportResult.SUCCESS, f"log export failed: {result}" assert tls_sink.received.get(timeout=5) == "/v1/logs" finally: exporter.shutdown() @@ -6639,7 +6645,7 @@ class TestOpenTelemetryProviderlessCallAttributes(unittest.TestCase): time.sleep(self.POLL_INTERVAL) return None - def _emitted_log_records(self, semconv_opt_in: str) -> tuple[LogData, ...]: + def _emitted_log_records(self, semconv_opt_in: str) -> tuple[ReadableLogRecord, ...]: log_exporter = InMemoryLogExporter() logger_provider = OTLoggerProvider() logger_provider.add_log_record_processor(SimpleLogRecordProcessor(log_exporter)) @@ -6713,7 +6719,12 @@ class TestDynamicTracerProviderCache(unittest.TestCase): must be bounded and must shut down whatever it drops (LIT-5437: threads accumulated until pods were OOMKilled).""" - BSP_THREAD_NAME = "OtelBatchSpanProcessor" + def setUp(self) -> None: + processor: Final = BatchSpanProcessor(InMemorySpanExporter()) + try: + self.bsp_thread_name = processor._batch_processor._worker_thread.name + finally: + processor.shutdown() def _logger(self, cap=3, exporter="console"): logger = OpenTelemetry( @@ -6730,7 +6741,7 @@ class TestDynamicTracerProviderCache(unittest.TestCase): logger._tracer_provider_cache.clear() def _live_exporter_threads(self): - return [t for t in threading.enumerate() if t.name == self.BSP_THREAD_NAME] + return [t for t in threading.enumerate() if t.name == self.bsp_thread_name] def _wait_for_exporter_threads(self, expected, timeout=10.0): """Dropped providers are shut down off-thread, so poll instead of sleeping.""" @@ -6751,8 +6762,6 @@ class TestDynamicTracerProviderCache(unittest.TestCase): logger._get_tracer_with_dynamic_headers({"authorization": f"Basic tenant-{i}"}) self.assertEqual(len(logger._tracer_provider_cache), 3) - # Guards the thread-name constant: a rename upstream would make this read 0 and the - # bound assertion below would pass while measuring nothing. self.assertGreaterEqual(len(self._live_exporter_threads()), 1) self.assertLessEqual(self._wait_for_exporter_threads(before + 3) - before, 3) diff --git a/tests/unit/interactions/conftest.py b/tests/unit/interactions/conftest.py new file mode 100644 index 00000000000..b53763c7036 --- /dev/null +++ b/tests/unit/interactions/conftest.py @@ -0,0 +1,9 @@ +import pytest + + +def pytest_addoption(parser: pytest.Parser) -> None: + parser.addoption( + "--interactions-spec-url", + default=None, + help="Run the Interactions OpenAPI compliance tests against this live spec instead of the vendored snapshot", + ) diff --git a/tests/unit/interactions/interactions.openapi.json b/tests/unit/interactions/interactions.openapi.json new file mode 100644 index 00000000000..3830f5093cb --- /dev/null +++ b/tests/unit/interactions/interactions.openapi.json @@ -0,0 +1,10024 @@ +{ + "openapi": "3.0.3", + "info": { + "title": "Gemini API", + "description": "The Gemini Interactions API allows developers to build generative AI applications using Gemini models. Gemini is our most capable model, built from the ground up to be multimodal. It can generalize and seamlessly understand, operate across, and combine different types of information including language, images, audio, video, and code. You can use the Gemini API for use cases like reasoning across text and images, content generation, dialogue agents, summarization and classification systems, and more.", + "version": "v1beta", + "x-google-revision": "0" + }, + "servers": [ + { + "url": "https://generativelanguage.googleapis.com", + "description": "Global Endpoint" + } + ], + "paths": { + "/upload/{api_version}/environments/{environment}/files/{path}": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + }, + { + "name": "environment", + "in": "path", + "description": "The ID of the environment that owns the destination file.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "path", + "in": "path", + "description": "The relative destination path inside the environment workspace.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "put": { + "operationId": "StartEnvironmentFileUpload", + "description": "Starts a resumable upload session for a file in an environment workspace.\nUpload the file bytes to the URL returned in the `X-Goog-Upload-URL`\nresponse header, using the resumable upload protocol.\n", + "parameters": [ + { + "$ref": "#/components/parameters/EnvironmentFileExtract" + }, + { + "$ref": "#/components/parameters/EnvironmentFileOverwrite" + }, + { + "$ref": "#/components/parameters/EnvironmentFileUploadCommand" + }, + { + "$ref": "#/components/parameters/EnvironmentFileUploadContentLength" + }, + { + "$ref": "#/components/parameters/EnvironmentFileUploadContentType" + }, + { + "$ref": "#/components/parameters/EnvironmentFileUploadProtocol" + } + ], + "responses": { + "200": { + "description": "Upload session created successfully. Send the file bytes to the returned upload URL.", + "headers": { + "X-Goog-Upload-URL": { + "$ref": "#/components/headers/EnvironmentFileUploadUrl" + } + } + }, + "4XX": { + "description": "The upload session request failed with a client error." + }, + "5XX": { + "description": "The upload session request failed with a server error." + } + }, + "summary": "Start an environment file upload", + "x-codeSamples": [ + { + "label": "start_upload", + "lang": "sh", + "source": "curl -i -X PUT \\\n 'https://generativelanguage.googleapis.com/upload/v1beta/environments/env_abc123/files/main.py?overwrite=true' \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H 'X-Goog-Upload-Protocol: resumable' \\\n -H 'X-Goog-Upload-Command: start' \\\n -H \"X-Goog-Upload-Header-Content-Length: $(wc -c < main.py)\" \\\n -H 'X-Goog-Upload-Header-Content-Type: text/x-python'\n" + } + ] + } + }, + "/{api_version}/agents": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "get": { + "operationId": "ListAgents", + "description": "Lists all Agents.", + "parameters": [ + { + "name": "page_size", + "in": "query", + "schema": { + "type": "integer", + "format": "int32" + } + }, + { + "name": "page_token", + "in": "query", + "schema": { + "type": "string" + } + }, + { + "name": "parent", + "in": "query", + "description": "Required. The parent resource to list agents from.\nFormat: `projects/{project}/locations/{location}`", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "examples": { + "list": { + "summary": "List Agents", + "value": { + "agents": [ + { + "description": "A helpful research assistant.", + "base_agent": "antigravity-preview-05-2026", + "id": "research-assistant-abc123", + "system_instruction": "You are a helpful research assistant.", + "tools": [ + { + "type": "google_search" + } + ] + } + ] + } + } + }, + "schema": { + "$ref": "#/components/schemas/ListAgentsResponse" + } + } + } + } + }, + "tags": [ + "agents" + ], + "x-codeSamples": [ + { + "label": "list", + "lang": "sh", + "source": "curl -X GET https://generativelanguage.googleapis.com/v1beta/agents \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "list", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.agents.list()\nfor agent in response.agents or []:\n print(agent.id)\n" + }, + { + "label": "list", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst agents = await ai.agents.list();\nfor (const agent of agents.agents ?? []) {\n console.log(agent.id);\n}\n" + }, + { + "label": "list", + "lang": "java", + "source": "import com.google.genai.gaos.models.agents.Agent;\nimport com.google.genai.gaos.models.agents.AgentListResponse;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.ListAgentsResponse;\n\nClient client = new Client();\nListAgentsResponse response = client.agents.list().call();\nList agents =\n response.agentListResponse().flatMap(AgentListResponse::agents).orElse(List.of());\nfor (Agent agent : agents) {\n agent.id().ifPresent(System.out::println);\n}\n" + }, + { + "label": "list", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n res, err := client.Agents.List(ctx, operations.ListAgentsRequest{})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.AgentListResponse != nil {\n for _, agent := range res.AgentListResponse.Agents {\n if agent.ID != nil {\n fmt.Println(*agent.ID)\n }\n }\n }\n}\n" + } + ] + }, + "post": { + "operationId": "CreateAgent", + "description": "Creates a new Agent (Typed version for SDK).", + "parameters": [ + { + "name": "parent", + "in": "query", + "description": "Required. The parent resource where this agent will be created.\nFormat: `projects/{project}/locations/{location}`", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "description": "Required. The request body.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Agent" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "examples": { + "create": { + "summary": "Create Agent", + "value": { + "description": "A helpful research assistant.", + "base_agent": "antigravity-preview-05-2026", + "base_environment": "remote", + "id": "research-assistant-abc123", + "system_instruction": "You are a helpful research assistant.", + "tools": [ + { + "type": "google_search" + } + ] + } + }, + "fork_from_env": { + "summary": "Agent Forked from Environment", + "value": { + "base_agent": "antigravity-preview-05-2026", + "base_environment": "env_abc123", + "id": "my-data-analyst-abc123", + "system_instruction": "You are a data analyst. Use the template at /workspace/template.py for all reports." + } + }, + "with_sources": { + "summary": "Agent with Sources", + "value": { + "base_agent": "antigravity-preview-05-2026", + "base_environment": { + "type": "remote", + "sources": [ + { + "type": "inline", + "content": "Always use matplotlib for charts. Include a summary table in every report.", + "target": ".agents/AGENTS.md" + }, + { + "type": "repository", + "source": "https://github.com/my-org/analysis-templates", + "target": "/workspace/templates" + } + ] + }, + "id": "data-analyst-abc123", + "system_instruction": "You are a data analyst. Always include visualizations and export results as PDF." + } + } + }, + "schema": { + "$ref": "#/components/schemas/Agent" + } + } + } + } + }, + "tags": [ + "agents" + ], + "x-codeSamples": [ + { + "label": "create", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/agents \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"id\": \"research-assistant-abc123\",\n \"base_agent\": \"antigravity-preview-05-2026\",\n \"description\": \"A helpful research assistant.\",\n \"system_instruction\": \"You are a helpful research assistant.\",\n \"base_environment\": \"remote\",\n \"tools\": [{\"type\": \"google_search\"}]\n }'\n" + }, + { + "label": "with_sources", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/agents \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"id\": \"data-analyst-abc123\",\n \"base_agent\": \"antigravity-preview-05-2026\",\n \"system_instruction\": \"You are a data analyst. Always include visualizations and export results as PDF.\",\n \"base_environment\": {\n \"type\": \"remote\",\n \"sources\": [\n {\n \"type\": \"inline\",\n \"target\": \".agents/AGENTS.md\",\n \"content\": \"Always use matplotlib for charts. Include a summary table in every report.\"\n },\n {\n \"type\": \"repository\",\n \"source\": \"https://github.com/my-org/analysis-templates\",\n \"target\": \"/workspace/templates\"\n }\n ]\n }\n }'\n" + }, + { + "label": "fork_from_env", + "lang": "sh", + "source": "# Step 1: Set up the environment interactively\nRESPONSE=$(curl -s -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"antigravity-preview-05-2026\",\n \"input\": \"Write a basic Hello World template to /workspace/template.py.\",\n \"environment\": \"remote\"\n }')\nENV_ID=$(echo \"$RESPONSE\" | python3 -c \"import sys,json; print(json.load(sys.stdin)['environment_id'])\")\n\n# Step 2: Fork that environment into a named agent\ncurl -X POST https://generativelanguage.googleapis.com/v1beta/agents \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d \"{\n \\\"id\\\": \\\"my-data-analyst\\\",\n \\\"base_agent\\\": \\\"antigravity-preview-05-2026\\\",\n \\\"system_instruction\\\": \\\"You are a data analyst. Use the template at /workspace/template.py for all reports.\\\",\n \\\"base_environment\\\": \\\"$ENV_ID\\\"\n }\"\n" + }, + { + "label": "create", + "lang": "python", + "source": "import uuid\nfrom google import genai\n\nclient = genai.Client()\nagent = client.agents.create(\n id=f\"research-assistant-{uuid.uuid4().hex[:8]}\",\n base_agent=\"antigravity-preview-05-2026\",\n description=\"A helpful research assistant.\",\n system_instruction=\"You are a helpful research assistant.\",\n base_environment=\"remote\",\n tools=[{\"type\": \"google_search\"}],\n)\nprint(agent.id)\n" + }, + { + "label": "with_sources", + "lang": "python", + "source": "import uuid\nfrom google import genai\n\nclient = genai.Client()\nagent = client.agents.create(\n id=f\"data-analyst-{uuid.uuid4().hex[:8]}\",\n base_agent=\"antigravity-preview-05-2026\",\n system_instruction=(\n \"You are a data analyst. Always include visualizations and export\"\n \" results as PDF.\"\n ),\n base_environment={\n \"type\": \"remote\",\n \"sources\": [\n {\n \"type\": \"inline\",\n \"target\": \".agents/AGENTS.md\",\n \"content\": (\n \"Always use matplotlib for charts. Include a summary\"\n \" table in every report.\"\n ),\n },\n {\n \"type\": \"repository\",\n \"source\": \"https://github.com/my-org/analysis-templates\",\n \"target\": \"/workspace/templates\",\n },\n ],\n },\n)\nprint(f\"Created agent: {agent.id}\")\n" + }, + { + "label": "fork_from_env", + "lang": "python", + "source": "import uuid\nfrom google import genai\n\nclient = genai.Client()\n\n# Step 1: Set up the environment interactively.\ninteraction = client.interactions.create(\n agent=\"antigravity-preview-05-2026\",\n input=\"Write a basic Hello World template to /workspace/template.py.\",\n environment=\"remote\",\n)\n\n# Step 2: Fork that environment into a named agent.\nagent_id = f\"my-data-analyst-{uuid.uuid4().hex[:8]}\"\nagent = client.agents.create(\n id=agent_id,\n base_agent=\"antigravity-preview-05-2026\",\n system_instruction=(\n \"You are a data analyst. Use the template at /workspace/template.py\"\n \" for all reports.\"\n ),\n base_environment=interaction.environment_id,\n)\nprint(f\"Forked agent: {agent.id}\")\n" + }, + { + "label": "create", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst agentId = `research-assistant-${crypto.randomUUID().slice(0, 8)}`;\nconst agent = await ai.agents.create({\n id: agentId,\n base_agent: 'antigravity-preview-05-2026',\n description: 'A helpful research assistant.',\n system_instruction: 'You are a helpful research assistant.',\n base_environment: 'remote',\n tools: [{type: 'google_search'}],\n});\nif (!agent.id) {\n throw new Error('Agent creation failed: ID is undefined');\n}\nconsole.log(agent.id);\n" + }, + { + "label": "with_sources", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst agentId = `data-analyst-${crypto.randomUUID().slice(0, 8)}`;\nconst agent = await ai.agents.create({\n id: agentId,\n base_agent: 'antigravity-preview-05-2026',\n system_instruction:\n 'You are a data analyst. Always include visualizations and export results as PDF.',\n base_environment: {\n type: 'remote',\n sources: [\n {\n type: 'inline',\n target: '.agents/AGENTS.md',\n content:\n 'Always use matplotlib for charts. Include a summary table in every report.',\n },\n {\n type: 'repository',\n source: 'https://github.com/my-org/analysis-templates',\n target: '/workspace/templates',\n },\n ],\n },\n});\nconsole.log(`Created agent: ${agent.id}`);\n" + }, + { + "label": "fork_from_env", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n// Step 1: Set up the environment interactively.\nconst interaction = await ai.interactions.create({\n agent: 'antigravity-preview-05-2026',\n input: 'Write a basic Hello World template to /workspace/template.py.',\n environment: 'remote',\n});\n\n// Step 2: Fork that environment into a named agent.\nconst agentId = `my-data-analyst-${crypto.randomUUID().slice(0, 8)}`;\nconst agent = await ai.agents.create({\n id: agentId,\n base_agent: 'antigravity-preview-05-2026',\n system_instruction:\n 'You are a data analyst. Use the template at /workspace/template.py for all reports.',\n base_environment: interaction.environment_id,\n});\nconsole.log(`Forked agent: ${agent.id}`);\n" + }, + { + "label": "create", + "lang": "java", + "source": "import com.google.genai.gaos.models.agents.Agent;\nimport com.google.genai.gaos.models.agents.BaseEnvironment;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateAgentResponse;\nimport com.google.genai.gaos.models.interactions.GoogleSearch;\n\nClient client = new Client();\nString agentId = \"research-assistant-\" + UUID.randomUUID().toString().substring(0, 8);\nAgent agent =\n Agent.builder()\n .id(agentId)\n .baseAgent(\"antigravity-preview-05-2026\")\n .description(\"A helpful research assistant.\")\n .systemInstruction(\"You are a helpful research assistant.\")\n .baseEnvironment(BaseEnvironment.of(\"remote\"))\n .tools(List.of(new GoogleSearch()))\n .build();\nCreateAgentResponse response = client.agents.create(agent);\nSystem.out.println(response.agent().flatMap(Agent::id).orElse(\"\"));\n" + }, + { + "label": "with_sources", + "lang": "java", + "source": "import com.google.genai.gaos.models.agents.Agent;\nimport com.google.genai.gaos.models.agents.BaseEnvironment;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateAgentResponse;\nimport com.google.genai.gaos.models.environments.Environment;\nimport com.google.genai.gaos.models.interactions.Source;\nimport com.google.genai.gaos.models.interactions.SourceType;\n\nClient client = new Client();\nString agentId = \"data-analyst-\" + UUID.randomUUID().toString().substring(0, 8);\ncom.google.genai.gaos.models.interactions.Environment env =\n com.google.genai.gaos.models.interactions.Environment.builder()\n .sources(\n List.of(\n Source.builder()\n .type(SourceType.INLINE)\n .target(\".agents/AGENTS.md\")\n .content(\n \"Always use matplotlib for charts. Include a summary table in every\"\n + \" report.\")\n .build(),\n Source.builder()\n .type(SourceType.REPOSITORY)\n .source(\"https://github.com/my-org/analysis-templates\")\n .target(\"/workspace/templates\")\n .build()))\n .build();\nAgent agent =\n Agent.builder()\n .id(agentId)\n .baseAgent(\"antigravity-preview-05-2026\")\n .systemInstruction(\n \"You are a data analyst. Always include visualizations and export results as PDF.\")\n .baseEnvironment(BaseEnvironment.of(env))\n .build();\nCreateAgentResponse response = client.agents.create(agent);\nSystem.out.println(\"Created agent: \" + response.agent().flatMap(Agent::id).orElse(\"\"));\n" + }, + { + "label": "fork_from_env", + "lang": "java", + "source": "import com.google.genai.gaos.models.agents.Agent;\nimport com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.gaos.models.agents.BaseEnvironment;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteractionEnvironment;\nimport com.google.genai.gaos.models.operations.CreateAgentResponse;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.Step;\n\nClient client = new Client();\n\n// Step 1: Set up the environment interactively.\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"antigravity-preview-05-2026\"))\n .input(\n InteractionsInput.of(\n \"Write a basic Hello World template to /workspace/template.py.\"))\n .environment(CreateAgentInteractionEnvironment.of(\"remote\"))\n .build();\nCreateInteractionResponse interactionResponse =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n interactionResponse\n .interaction()\n .orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\n\n// Step 2: Fork that environment into a named agent.\nString agentId = \"my-data-analyst-\" + UUID.randomUUID().toString().substring(0, 8);\nAgent agent =\n Agent.builder()\n .id(agentId)\n .baseAgent(\"antigravity-preview-05-2026\")\n .systemInstruction(\n \"You are a data analyst. Use the template at /workspace/template.py for all\"\n + \" reports.\")\n .baseEnvironment(BaseEnvironment.of(interaction.environmentId().orElse(\"\")))\n .build();\nCreateAgentResponse agentResponse = client.agents.create(agent);\nSystem.out.println(\"Forked agent: \" + agentResponse.agent().flatMap(Agent::id).orElse(\"\"));\n" + }, + { + "label": "create", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"time\"\n \"google.golang.org/genai\"\n gaos_agents \"google.golang.org/genai/interactions/models/agents\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n agentID := fmt.Sprintf(\"research-assistant-%d\", time.Now().UnixNano())\n req := operations.CreateAgentRequest{\n Body: gaos_agents.Agent{\n ID: genai.Ptr(agentID),\n BaseAgent: genai.Ptr(\"antigravity-preview-05-2026\"),\n Description: genai.Ptr(\"A helpful research assistant.\"),\n SystemInstruction: genai.Ptr(\"You are a helpful research assistant.\"),\n BaseEnvironment: genai.Ptr(gaos_agents.NewBaseEnvironment(\"remote\")),\n Tools: []gaos_agents.AgentTool{\n gaos_agents.NewAgentTool(gaos_interactions.GoogleSearch{}),\n },\n },\n }\n\n res, err := client.Agents.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Agent != nil && res.Agent.ID != nil {\n fmt.Println(*res.Agent.ID)\n }\n}\n" + }, + { + "label": "with_sources", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"time\"\n \"google.golang.org/genai\"\n gaos_agents \"google.golang.org/genai/interactions/models/agents\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n agentID := fmt.Sprintf(\"data-analyst-%d\", time.Now().UnixNano())\n env := gaos_interactions.Environment{\n Sources: []gaos_interactions.Source{\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeInline),\n Target: genai.Ptr(\".agents/AGENTS.md\"),\n Content: genai.Ptr(\"Always use matplotlib for charts. Include a summary table in every report.\"),\n },\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeRepository),\n Source: genai.Ptr(\"https://github.com/my-org/analysis-templates\"),\n Target: genai.Ptr(\"/workspace/templates\"),\n },\n },\n }\n\n req := operations.CreateAgentRequest{\n Body: gaos_agents.Agent{\n ID: genai.Ptr(agentID),\n BaseAgent: genai.Ptr(\"antigravity-preview-05-2026\"),\n SystemInstruction: genai.Ptr(\"You are a data analyst. Always include visualizations and export results as PDF.\"),\n BaseEnvironment: genai.Ptr(gaos_agents.NewBaseEnvironment(env)),\n },\n }\n\n res, err := client.Agents.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Agent != nil && res.Agent.ID != nil {\n fmt.Printf(\"Created agent: %s\\n\", *res.Agent.ID)\n }\n}\n" + }, + { + "label": "fork_from_env", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"time\"\n \"google.golang.org/genai\"\n gaos_agents \"google.golang.org/genai/interactions/models/agents\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n // Step 1: Set up the environment interactively.\n interactionReq := operations.CreateInteractionRequest{\n Body: operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"antigravity-preview-05-2026\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Write a basic Hello World template to /workspace/template.py.\")),\n Environment: genai.Ptr(gaos_interactions.NewCreateAgentInteractionEnvironment(\"remote\")),\n }),\n }\n interactionRes, err := client.Interactions.Create(ctx, interactionReq)\n if err != nil {\n log.Fatal(err)\n }\n envID := \"\"\n if interactionRes.Interaction != nil && interactionRes.Interaction.EnvironmentID != nil {\n envID = *interactionRes.Interaction.EnvironmentID\n }\n\n // Step 2: Fork that environment into a named agent.\n agentID := fmt.Sprintf(\"my-data-analyst-%d\", time.Now().UnixNano())\n req := operations.CreateAgentRequest{\n Body: gaos_agents.Agent{\n ID: genai.Ptr(agentID),\n BaseAgent: genai.Ptr(\"antigravity-preview-05-2026\"),\n SystemInstruction: genai.Ptr(\"You are a data analyst. Use the template at /workspace/template.py for all reports.\"),\n BaseEnvironment: genai.Ptr(gaos_agents.NewBaseEnvironment(envID)),\n },\n }\n res, err := client.Agents.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n if res.Agent != nil && res.Agent.ID != nil {\n fmt.Printf(\"Forked agent: %s\\n\", *res.Agent.ID)\n }\n}\n" + } + ] + } + }, + "/{api_version}/agents/{agentsId}": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "delete": { + "operationId": "DeleteAgent", + "description": "Deletes an Agent.", + "parameters": [ + { + "name": "agentsId", + "in": "path", + "description": "Required. The name of the agent to delete.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "examples": { + "delete": { + "summary": "Delete Agent", + "value": {} + } + }, + "schema": { + "$ref": "#/components/schemas/Empty" + } + } + } + } + }, + "tags": [ + "agents" + ], + "x-codeSamples": [ + { + "label": "delete", + "lang": "sh", + "source": "curl -X DELETE https://generativelanguage.googleapis.com/v1beta/agents/ag_abc123 \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "delete", + "lang": "python", + "source": "import uuid\nfrom google import genai\n\nclient = genai.Client()\n\n\nclient.agents.delete(agent_id)\nprint(\"Agent deleted successfully.\")\n" + }, + { + "label": "delete", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nawait ai.agents.delete(agentId);\nconsole.log('Agent deleted successfully.');\n" + }, + { + "label": "delete", + "lang": "java", + "source": "import com.google.genai.gaos.models.agents.Agent;\nimport com.google.genai.Client;\n\nClient client = new Client();\n\n\nclient.agents.delete(agentId);\nSystem.out.println(\"Agent deleted successfully.\");\n" + }, + { + "label": "delete", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n _, err = client.Agents.Delete(ctx, operations.DeleteAgentRequest{ID: agentID})\n if err != nil {\n log.Fatal(err)\n }\n fmt.Println(\"Agent deleted successfully.\")\n}\n" + } + ] + }, + "get": { + "operationId": "GetAgent", + "description": "Gets a specific Agent.", + "parameters": [ + { + "name": "agentsId", + "in": "path", + "description": "Required. The name of the agent to retrieve.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "examples": { + "get": { + "summary": "Get Agent", + "value": { + "created": "2025-11-26T12:25:15Z", + "display_name": "My Research Agent", + "id": "ag_abc123", + "object": "agent", + "system_instruction": "You are a helpful research assistant.", + "tools": [ + { + "type": "google_search" + } + ], + "updated": "2025-11-26T12:25:15Z" + } + } + }, + "schema": { + "$ref": "#/components/schemas/Agent" + } + } + } + } + }, + "tags": [ + "agents" + ], + "x-codeSamples": [ + { + "label": "get", + "lang": "sh", + "source": "curl -X GET https://generativelanguage.googleapis.com/v1beta/agents/ag_abc123 \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "get", + "lang": "python", + "source": "import uuid\nfrom google import genai\n\nclient = genai.Client()\n\n\nagent = client.agents.get(agent_id)\nprint(agent.id)\n" + }, + { + "label": "get", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nconst agent = await ai.agents.get(agentId);\nif (!agent.id) {\n throw new Error('Agent retrieval failed: ID is undefined');\n}\nconsole.log(agent.id);\n" + }, + { + "label": "get", + "lang": "java", + "source": "import com.google.genai.gaos.models.agents.Agent;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.GetAgentResponse;\n\nClient client = new Client();\n\n\nGetAgentResponse response = client.agents.get(agentId);\nAgent agent = response.agent().orElseThrow(() -> new RuntimeException(\"No agent returned\"));\nSystem.out.println(agent.id().orElse(\"\"));\n" + }, + { + "label": "get", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Agents.Get(ctx, operations.GetAgentRequest{ID: agentID})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Agent != nil && res.Agent.ID != nil {\n fmt.Println(*res.Agent.ID)\n }\n}\n" + } + ] + } + }, + "/{api_version}/credentials": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "get": { + "operationId": "ListCredentials", + "description": "Lists credentials.", + "parameters": [ + { + "name": "page_size", + "in": "query", + "description": "Optional. Maximum number of credentials to return.", + "schema": { + "type": "integer", + "format": "int32" + } + }, + { + "name": "page_token", + "in": "query", + "description": "Optional. Pagination token.", + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListCredentialsResponse" + } + } + } + } + }, + "tags": [ + "credentials" + ] + }, + "post": { + "operationId": "CreateCredential", + "description": "Creates a new credential.", + "parameters": [], + "requestBody": { + "description": "Required. The request body.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CredentialCreateParams" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Credential" + } + } + } + } + }, + "tags": [ + "credentials" + ] + } + }, + "/{api_version}/credentials/{id}": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "delete": { + "operationId": "DeleteCredential", + "description": "Deletes a credential.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. Resource ID segment making up resource `name`. It identifies the resource\nwithin its parent collection as described in https://google.aip.dev/122.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Empty" + } + } + } + } + }, + "tags": [ + "credentials" + ] + }, + "get": { + "operationId": "GetCredential", + "description": "Gets a credential by ID.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. Resource ID segment making up resource `name`. It identifies the resource\nwithin its parent collection as described in https://google.aip.dev/122.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Credential" + } + } + } + } + }, + "tags": [ + "credentials" + ] + }, + "patch": { + "operationId": "UpdateCredential", + "description": "Updates a credential.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. Resource ID segment making up resource `name`. It identifies the resource\nwithin its parent collection as described in https://google.aip.dev/122.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "update_mask", + "in": "query", + "description": "Optional. The list of fields to update.", + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "description": "Required. The request body.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CredentialUpdateParams" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Credential" + } + } + } + } + }, + "tags": [ + "credentials" + ] + } + }, + "/{api_version}/environments": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "get": { + "operationId": "ListEnvironments", + "description": "Lists environments.", + "parameters": [ + { + "name": "page_size", + "in": "query", + "description": "Optional. Maximum number of environments to return.\nIf unspecified, defaults to 50. Maximum is 1000.", + "schema": { + "type": "integer", + "format": "int32" + } + }, + { + "name": "page_token", + "in": "query", + "description": "Optional. Pagination token.", + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListEnvironmentsResponse" + } + } + } + } + }, + "tags": [ + "environments" + ], + "x-codeSamples": [ + { + "label": "list", + "lang": "sh", + "source": "curl -X GET https://generativelanguage.googleapis.com/v1beta/environments \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "list", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.environments.list()\nfor environment in response.environments or []:\n print(environment.id)\n" + }, + { + "label": "list", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst response = await ai.environments.list();\nfor (const env of response.environments ?? []) {\n console.log(env.id);\n}\n" + }, + { + "label": "list", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.environments.Environment;\nimport com.google.genai.gaos.models.operations.ListEnvironmentsResponse;\n\nClient client = new Client();\nListEnvironmentsResponse response = client.environments.listEnvironmentsDirect();\nfor (com.google.genai.gaos.models.environments.Environment env :\n response.listEnvironmentsResponse().flatMap(res -> res.environments()).orElse(List.of())) {\n System.out.println(env.id().orElse(\"\"));\n}\n" + }, + { + "label": "list", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n res, err := client.Environments.ListEnvironments(ctx, operations.ListEnvironmentsRequest{})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.ListEnvironmentsResponse != nil {\n for _, env := range res.ListEnvironmentsResponse.Environments {\n if env.ID != \"\" {\n fmt.Println(env.ID)\n }\n }\n }\n}\n" + } + ] + }, + "post": { + "operationId": "CreateEnvironment", + "description": "Creates an environment.", + "parameters": [], + "requestBody": { + "description": "Required. The request body.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CreateEnvironmentRequest" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Environment" + } + } + } + } + }, + "tags": [ + "environments" + ], + "x-codeSamples": [ + { + "label": "create", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/environments \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"sources\": [{\n \"type\": \"inline\",\n \"target\": \"main.py\",\n \"content\": \"print(\\\"Hello, World!\\\")\"\n }]\n }'\n" + }, + { + "label": "copy", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/environments \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"from_environment\": \"environments/env_abc123\"\n }'\n" + }, + { + "label": "create", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nenvironment = client.environments.create(\n sources=[{\n \"type\": \"inline\",\n \"target\": \"main.py\",\n \"content\": \"print('Hello, World!')\",\n }]\n)\nprint(environment.id)\n" + }, + { + "label": "copy", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n\nenvironment = client.environments.create(\n from_environment=source_env.id,\n)\nprint(environment.id)\n" + }, + { + "label": "create", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst environment = await ai.environments.create({\n sources: [\n {\n type: 'inline',\n target: 'main.py',\n content: \"print('Hello, World!')\",\n },\n ],\n});\nif (!environment.id) {\n throw new Error('Environment creation failed: ID is undefined');\n}\nconsole.log(environment.id);\n" + }, + { + "label": "copy", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nconst environment = await ai.environments.create({\n from_environment: sourceEnv.id,\n});\nif (!environment.id) {\n throw new Error('Environment creation failed: ID is undefined');\n}\nconsole.log(environment.id);\n" + }, + { + "label": "create", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.environments.CreateEnvironmentRequest;\nimport com.google.genai.gaos.models.environments.Environment;\nimport com.google.genai.gaos.models.interactions.Source;\nimport com.google.genai.gaos.models.interactions.SourceType;\n\nClient client = new Client();\nCreateEnvironmentRequest request =\n CreateEnvironmentRequest.builder()\n .sources(\n List.of(\n Source.builder()\n .type(SourceType.INLINE)\n .target(\"main.py\")\n .content(\"print('Hello, World!')\")\n .build()))\n .build();\ncom.google.genai.gaos.models.environments.Environment environment =\n client.environments.createEnvironment(request).environment().orElseThrow();\nSystem.out.println(environment.id().orElse(\"\"));\n" + }, + { + "label": "copy", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.environments.CreateEnvironmentRequest;\nimport com.google.genai.gaos.models.environments.Environment;\n\nClient client = new Client();\n\n\nCreateEnvironmentRequest request =\n CreateEnvironmentRequest.builder().fromEnvironment(sourceEnvId).build();\ncom.google.genai.gaos.models.environments.Environment environment =\n client.environments.createEnvironment(request).environment().orElseThrow();\nSystem.out.println(environment.id().orElse(\"\"));\n" + }, + { + "label": "create", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n req := operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n Sources: []gaos_interactions.Source{\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeInline),\n Target: genai.Ptr(\"main.py\"),\n Content: genai.Ptr(\"print('Hello, World!')\"),\n },\n },\n },\n }\n\n res, err := client.Environments.CreateEnvironment(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" + }, + { + "label": "copy", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.CreateEnvironment(ctx, operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n FromEnvironment: genai.Ptr(sourceEnvID),\n },\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" + } + ] + } + }, + "/{api_version}/environments/{environment}/files/{path}": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + }, + { + "name": "environment", + "in": "path", + "description": "The ID of the environment whose snapshot to read.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "path", + "in": "path", + "description": "Path of the file or directory inside the environment workspace, relative to its root (e.g. src).", + "required": true, + "schema": { + "type": "string" + } + } + ], + "get": { + "operationId": "GetEnvironmentFiles", + "description": "Retrieves file metadata or directory contents from an environment's snapshot. To download file content, use the download URL returned in the response.", + "parameters": [ + { + "name": "page_size", + "in": "query", + "description": "Optional. Maximum number of entries to return per page (for directory listing).", + "schema": { + "type": "integer", + "format": "int32" + } + }, + { + "name": "page_token", + "in": "query", + "description": "Optional. Pagination token for directory listing.", + "schema": { + "type": "string" + } + }, + { + "name": "recursive", + "in": "query", + "description": "Optional. If true and the path is a directory, recursively lists all files.", + "schema": { + "type": "boolean" + } + } + ], + "responses": { + "200": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/GetEnvironmentFilesResponse" + } + } + } + } + }, + "tags": [ + "environments.files" + ], + "x-codeSamples": [ + { + "label": "list_files", + "lang": "sh", + "source": "curl -X GET 'https://generativelanguage.googleapis.com/v1beta/environments/env_abc123/files/src' \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "get_file", + "lang": "sh", + "source": "curl -X GET 'https://generativelanguage.googleapis.com/v1beta/environments/env_abc123/files/main.py' \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "download_file", + "lang": "sh", + "source": "curl -X GET 'https://generativelanguage.googleapis.com/v1beta/environments/env_abc123/files/src/main.py?alt=media' \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n --output main.py\n" + }, + { + "label": "list_files", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n\nresponse = client.environments.files.list(\n environment=created.id,\n path=\"src\",\n)\nfor file in response.files or []:\n print(file.name, file.type, file.size_bytes)\n" + }, + { + "label": "get_file", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n\nresponse = client.environments.files.get(\n environment=created.id,\n path=\"main.py\",\n)\nfor file in response.files or []:\n print(file.name, file.size_bytes)\n" + }, + { + "label": "download_file", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n\ncontent = client.environments.files.download(\n environment=created.id,\n path=\"src/main.py\",\n)\nprint(content.decode(\"utf-8\"))\n" + }, + { + "label": "list_files", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nconst response = await ai.environments.files.list({\n environment: created.id,\n path: 'src',\n});\nfor (const file of response.files ?? []) {\n console.log(file.name, file.type, file.size_bytes);\n}\n" + }, + { + "label": "get_file", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nconst response = await ai.environments.files.list({\n environment: created.id,\n path: 'main.py',\n});\nfor (const file of response.files ?? []) {\n console.log(file.name, file.size_bytes);\n}\n" + }, + { + "label": "list_files", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.environments.EnvironmentFile;\nimport com.google.genai.gaos.models.operations.GetEnvironmentFilesRequest;\nimport com.google.genai.gaos.models.environments.GetEnvironmentFilesResponse;\n\nClient client = new Client();\n\n\ncom.google.genai.gaos.models.operations.GetEnvironmentFilesResponse response =\n client\n .environments\n .files()\n .list(GetEnvironmentFilesRequest.builder().environment(envId).path(\"src\").build());\nGetEnvironmentFilesResponse filesResponse =\n response.getEnvironmentFilesResponse().orElseThrow();\nfor (EnvironmentFile file : filesResponse.files().orElse(List.of())) {\n System.out.println(\n file.name().orElse(\"\")\n + \" \"\n + file.type().map(Object::toString).orElse(\"\")\n + \" \"\n + file.sizeBytes().orElse(\"\"));\n}\n" + }, + { + "label": "get_file", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.environments.EnvironmentFile;\nimport com.google.genai.gaos.models.operations.GetEnvironmentFilesRequest;\nimport com.google.genai.gaos.models.environments.GetEnvironmentFilesResponse;\n\nClient client = new Client();\n\n\ncom.google.genai.gaos.models.operations.GetEnvironmentFilesResponse response =\n client\n .environments\n .files()\n .list(GetEnvironmentFilesRequest.builder().environment(envId).path(\"main.py\").build());\nGetEnvironmentFilesResponse filesResponse =\n response.getEnvironmentFilesResponse().orElseThrow();\nfor (EnvironmentFile file : filesResponse.files().orElse(List.of())) {\n System.out.println(file.name().orElse(\"\") + \" \" + file.sizeBytes().orElse(\"\"));\n}\n" + }, + { + "label": "list_files", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"src\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n fileType := \"\"\n if file.Type != nil {\n fileType = string(*file.Type)\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, fileType, size)\n }\n }\n}\n" + }, + { + "label": "get_file", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"main.py\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, size)\n }\n }\n}\n" + } + ] + } + }, + "/{api_version}/environments/{id}": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "delete": { + "operationId": "DeleteEnvironment", + "description": "Deletes an environment.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. Resource ID segment making up resource `name`. It identifies the resource\nwithin its parent collection as described in https://google.aip.dev/122.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Empty" + } + } + } + } + }, + "tags": [ + "environments" + ], + "x-codeSamples": [ + { + "label": "delete", + "lang": "sh", + "source": "curl -X DELETE https://generativelanguage.googleapis.com/v1beta/environments/env_abc123 \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "delete", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n\nclient.environments.delete(id=created.id)\nprint(\"Environment deleted successfully.\")\n" + }, + { + "label": "delete", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nawait ai.environments.delete(created.id);\nconsole.log('Environment deleted successfully.');\n" + }, + { + "label": "delete", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.environments.Environment;\n\nClient client = new Client();\n\n\nclient.environments.deleteEnvironment(envId);\nSystem.out.println(\n \"com.google.genai.gaos.models.environments.Environment deleted successfully.\");\n" + }, + { + "label": "delete", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n _, err = client.Environments.DeleteEnvironment(ctx, operations.DeleteEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n fmt.Println(\"Environment deleted successfully.\")\n}\n" + } + ] + }, + "get": { + "operationId": "GetEnvironment", + "description": "Gets an environment.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. Resource ID segment making up resource `name`. It identifies the resource\nwithin its parent collection as described in https://google.aip.dev/122.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Environment" + } + } + } + } + }, + "tags": [ + "environments" + ], + "x-codeSamples": [ + { + "label": "get", + "lang": "sh", + "source": "curl -X GET https://generativelanguage.googleapis.com/v1beta/environments/env_abc123 \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "get", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n\nenvironment = client.environments.get(id=created.id)\nprint(environment.id)\n" + }, + { + "label": "get", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nconst environment = await ai.environments.get(created.id);\nconsole.log(environment.id);\n" + }, + { + "label": "get", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.environments.Environment;\n\nClient client = new Client();\n\n\ncom.google.genai.gaos.models.environments.Environment environment =\n client.environments.getEnvironment(envId).environment().orElseThrow();\nSystem.out.println(environment.id().orElse(\"\"));\n" + }, + { + "label": "get", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.GetEnvironment(ctx, operations.GetEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n }\n}\n" + } + ] + } + }, + "/{api_version}/interactions": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "post": { + "operationId": "CreateInteraction", + "description": "Creates a new interaction.", + "requestBody": { + "description": "The request body.", + "required": true, + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "$ref": "#/components/schemas/CreateAgentInteractionParams" + }, + { + "$ref": "#/components/schemas/CreateModelInteractionParams" + } + ] + } + } + } + }, + "responses": { + "200": { + "description": "Successful operation", + "content": { + "application/json": { + "examples": { + "antigravity": { + "summary": "Antigravity Agent", + "value": { + "agent": "antigravity-preview-05-2026", + "created": "2025-11-26T12:22:47Z", + "environment_id": "env_abc123", + "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", + "object": "interaction", + "status": "completed", + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "I've summarized the top 5 Hacker News stories and saved the results to /workspace/summary.md." + } + ] + } + ], + "updated": "2025-11-26T12:22:47Z", + "usage": { + "input_tokens_by_modality": [ + { + "modality": "text", + "tokens": 50 + } + ], + "total_cached_tokens": 0, + "total_input_tokens": 50, + "total_output_tokens": 500, + "total_thought_tokens": 200, + "total_tokens": 750, + "total_tool_use_tokens": 0 + } + } + }, + "deep_research": { + "summary": "Deep Research", + "value": { + "agent": "deep-research-pro-preview-12-2025", + "created": "2025-11-26T12:22:47Z", + "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", + "object": "interaction", + "status": "completed", + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "Here is a comprehensive research report on the current state of cancer research..." + } + ] + } + ], + "updated": "2025-11-26T12:22:47Z", + "usage": { + "input_tokens_by_modality": [ + { + "modality": "text", + "tokens": 20 + } + ], + "total_cached_tokens": 0, + "total_input_tokens": 20, + "total_output_tokens": 1000, + "total_thought_tokens": 500, + "total_tokens": 1520, + "total_tool_use_tokens": 0 + } + } + }, + "function_calling": { + "summary": "Function Calling", + "value": { + "created": "2025-11-26T12:22:47Z", + "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", + "model": "gemini-3.6-flash", + "object": "interaction", + "status": "requires_action", + "steps": [ + { + "name": "get_weather", + "type": "function_call", + "arguments": { + "location": "Boston, MA" + }, + "id": "gth23981" + } + ], + "updated": "2025-11-26T12:22:47Z", + "usage": { + "input_tokens_by_modality": [ + { + "modality": "text", + "tokens": 100 + } + ], + "total_cached_tokens": 0, + "total_input_tokens": 100, + "total_output_tokens": 25, + "total_thought_tokens": 0, + "total_tokens": 125, + "total_tool_use_tokens": 50 + } + } + }, + "multi_turn": { + "summary": "Multi-turn", + "value": { + "created": "2025-11-26T12:22:47Z", + "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", + "model": "gemini-3.6-flash", + "object": "interaction", + "status": "completed", + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris." + } + ] + } + ], + "updated": "2025-11-26T12:22:47Z", + "usage": { + "input_tokens_by_modality": [ + { + "modality": "text", + "tokens": 50 + } + ], + "total_cached_tokens": 0, + "total_input_tokens": 50, + "total_output_tokens": 10, + "total_thought_tokens": 0, + "total_tokens": 60, + "total_tool_use_tokens": 0 + } + } + }, + "multimodal_image": { + "summary": "Image Input", + "value": { + "created": "2025-11-26T12:22:47Z", + "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", + "model": "gemini-3.6-flash", + "object": "interaction", + "status": "completed", + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "A white humanoid robot with glowing blue eyes stands holding a red skateboard." + } + ] + } + ], + "updated": "2025-11-26T12:22:47Z", + "usage": { + "input_tokens_by_modality": [ + { + "modality": "text", + "tokens": 10 + }, + { + "modality": "image", + "tokens": 258 + } + ], + "total_cached_tokens": 0, + "total_input_tokens": 268, + "total_output_tokens": 20, + "total_thought_tokens": 0, + "total_tokens": 288, + "total_tool_use_tokens": 0 + } + } + }, + "reuse_env": { + "summary": "Reuse Environment", + "value": { + "agent": "antigravity-preview-05-2026", + "created": "2025-11-26T12:23:00Z", + "environment_id": "env_abc123", + "id": "v1_Chd2ZTJhYmNkZWZnaGlqa2xtbm9wcXJzdHV2d3h5ejAxMjM0NTY3ODkwMTIzNDU2Nzg", + "object": "interaction", + "status": "completed", + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "I've updated /workspace/hello.py to accept a name argument and greet the user." + } + ] + } + ], + "updated": "2025-11-26T12:23:00Z", + "usage": { + "input_tokens_by_modality": [ + { + "modality": "text", + "tokens": 80 + } + ], + "total_cached_tokens": 0, + "total_input_tokens": 80, + "total_output_tokens": 200, + "total_thought_tokens": 100, + "total_tokens": 380, + "total_tool_use_tokens": 0 + } + } + }, + "simple": { + "summary": "Simple Request", + "value": { + "created": "2025-11-26T12:25:15Z", + "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", + "model": "gemini-3.6-flash", + "object": "interaction", + "status": "completed", + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "Hello! I'm functioning perfectly and ready to assist you.\n\nHow are you doing today?" + } + ] + } + ], + "updated": "2025-11-26T12:25:15Z", + "usage": { + "input_tokens_by_modality": [ + { + "modality": "text", + "tokens": 7 + } + ], + "total_cached_tokens": 0, + "total_input_tokens": 7, + "total_output_tokens": 20, + "total_thought_tokens": 22, + "total_tokens": 49, + "total_tool_use_tokens": 0 + } + } + } + }, + "schema": { + "$ref": "#/components/schemas/Interaction" + } + }, + "text/event-stream": { + "schema": { + "$ref": "#/components/schemas/InteractionSseStreamEvent" + }, + "x-speakeasy-sse-sentinel": "[DONE]" + } + } + }, + "4XX": { + "description": "Error creating interaction", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "$ref": "#/components/schemas/Error" + } + }, + "required": [ + "error" + ] + } + } + } + }, + "5XX": { + "description": "Error creating interaction", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "$ref": "#/components/schemas/Error" + } + }, + "required": [ + "error" + ] + } + } + } + } + }, + "summary": "Creating an interaction", + "x-speakeasy-exports": [ + { + "name": "InteractionCreateParams", + "group": "interactions", + "representation": "input" + } + ], + "x-speakeasy-group": "interactions", + "x-speakeasy-name-override": "create", + "x-codeSamples": [ + { + "label": "simple", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"input\": \"Hello, how are you?\"\n }'\n" + }, + { + "label": "multi_turn", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"input\": [\n { \"type\": \"user_input\", \"content\": [{ \"type\": \"text\", \"text\": \"Hello!\" }] },\n { \"type\": \"model_output\", \"content\": [{ \"type\": \"text\", \"text\": \"Hi there! How can I help you today?\" }] },\n { \"type\": \"user_input\", \"content\": [{ \"type\": \"text\", \"text\": \"What is the capital of France?\" }] }\n ]\n }'\n" + }, + { + "label": "multimodal_image", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"input\": [\n {\n \"type\": \"text\",\n \"text\": \"What is in this picture?\"\n },\n {\n \"type\": \"image\",\n \"data\": \"BASE64_ENCODED_IMAGE\",\n \"mime_type\": \"image/png\"\n }\n ]\n }'\n" + }, + { + "label": "function_calling", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"tools\": [\n {\n \"type\": \"function\",\n \"name\": \"get_weather\",\n \"description\": \"Get the current weather in a given location\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"location\": {\n \"type\": \"string\",\n \"description\": \"The city and state, e.g. San Francisco, CA\"\n }\n },\n \"required\": [\n \"location\"\n ]\n }\n }\n ],\n \"input\": \"What is the weather like in Boston, MA?\"\n }'\n" + }, + { + "label": "deep_research", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"deep-research-pro-preview-12-2025\",\n \"input\": \"Find a cure to cancer\",\n \"background\": true\n }'\n" + }, + { + "label": "antigravity", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"antigravity-preview-05-2026\",\n \"input\": \"Read Hacker News, summarize the top 5 stories, and save results as a markdown file.\",\n \"environment\": \"remote\"\n }'\n" + }, + { + "label": "reuse_env", + "lang": "sh", + "source": "# Step 1: Create an interaction with a fresh remote environment.\nRESPONSE=$(curl -s -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -H \"Api-Revision: 2026-05-20\" \\\n -d '{\n \"agent\": \"antigravity-preview-05-2026\",\n \"input\": \"Write a hello world script at /workspace/hello.py.\",\n \"environment\": \"remote\"\n }')\nINTERACTION_ID=$(echo \"$RESPONSE\" | python3 -c \"import sys,json; print(json.load(sys.stdin)['id'])\")\nENV_ID=$(echo \"$RESPONSE\" | python3 -c \"import sys,json; print(json.load(sys.stdin)['environment_id'])\")\n\n# Step 2: Reuse the same environment in a follow-up interaction.\ncurl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d \"{\n \\\"agent\\\": \\\"antigravity-preview-05-2026\\\",\n \\\"input\\\": \\\"Modify the script to accept a name argument and greet the user.\\\",\n \\\"environment\\\": \\\"$ENV_ID\\\",\n \\\"previous_interaction_id\\\": \\\"$INTERACTION_ID\\\"\n }\"\n" + }, + { + "label": "with_sources", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"antigravity-preview-05-2026\",\n \"input\": \"List all files under /workspace and summarize what you find.\",\n \"environment\": {\n \"type\": \"remote\",\n \"sources\": [\n {\n \"type\": \"repository\",\n \"source\": \"https://github.com/octocat/Spoon-Knife\",\n \"target\": \"/workspace/repo\"\n },\n {\n \"type\": \"inline\",\n \"content\": \"Focus on Python files only.\",\n \"target\": \"/workspace/notes.txt\"\n }\n ]\n }\n }'\n" + }, + { + "label": "custom_agent", + "lang": "sh", + "source": "# Step 1: Create a custom agent.\ncurl -X POST https://generativelanguage.googleapis.com/v1beta/agents \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"id\": \"code-reviewer\",\n \"base_agent\": \"antigravity-preview-05-2026\",\n \"system_instruction\": \"You are a senior code reviewer. Check every file for bugs, style issues, and security vulnerabilities.\",\n \"base_environment\": {\n \"type\": \"remote\",\n \"sources\": [{\n \"type\": \"repository\",\n \"source\": \"https://github.com/octocat/Spoon-Knife\",\n \"target\": \"/workspace/repo\"\n }]\n }\n }'\n\n# Step 2: Use the custom agent.\ncurl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"code-reviewer\",\n \"input\": \"Review the latest changes in /workspace/repo/src and file a summary.\",\n \"environment\": \"remote\"\n }'\n" + }, + { + "label": "simple", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\ninteraction = client.interactions.create(\n model=\"gemini-flash-latest\",\n input=\"Hello, how are you?\",\n)\nprint(interaction.output_text)\n" + }, + { + "label": "multi_turn", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n input=[\n {\n \"type\": \"user_input\",\n \"content\": [{\"type\": \"text\", \"text\": \"Hello!\"}],\n },\n {\n \"type\": \"model_output\",\n \"content\": [{\n \"type\": \"text\",\n \"text\": \"Hi there! How can I help you today?\",\n }],\n },\n {\n \"type\": \"user_input\",\n \"content\": [\n {\"type\": \"text\", \"text\": \"What is the capital of France?\"}\n ],\n },\n ],\n)\nprint(response.output_text)\n" + }, + { + "label": "multimodal_image", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n input=[\n {\"type\": \"text\", \"text\": \"What is in this picture?\"},\n {\n \"type\": \"image\",\n \"data\": \"BASE64_ENCODED_IMAGE\",\n \"mime_type\": \"image/png\",\n },\n ],\n)\nprint(response.output_text)\n" + }, + { + "label": "function_calling", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n tools=[{\n \"type\": \"function\",\n \"name\": \"get_weather\",\n \"description\": \"Get the current weather in a given location\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"location\": {\n \"type\": \"string\",\n \"description\": (\n \"The city and state, e.g. San Francisco, CA\"\n ),\n }\n },\n \"required\": [\"location\"],\n },\n }],\n input=\"What is the weather like in Boston, MA?\",\n)\nprint(response.steps[-1])\n" + }, + { + "label": "deep_research", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\ninteraction = client.interactions.create(\n agent=\"deep-research-pro-preview-12-2025\",\n input=\"find a cure to cancer\",\n background=True,\n)\nprint(interaction.status)\n" + }, + { + "label": "antigravity", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\ninteraction = client.interactions.create(\n agent=\"antigravity-preview-05-2026\",\n input=(\n \"Read Hacker News, summarize the top 5 stories, and save results as\"\n \" a markdown file.\"\n ),\n environment=\"remote\",\n)\nprint(interaction.output_text)\n" + }, + { + "label": "reuse_env", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n# Step 1: Create an interaction with a fresh remote environment.\ninteraction = client.interactions.create(\n agent=\"antigravity-preview-05-2026\",\n input=\"Write a hello world script at /workspace/hello.py.\",\n environment=\"remote\",\n)\nprint(f\"Environment ID: {interaction.environment_id}\")\n\n# Step 2: Reuse the same environment in a follow-up interaction.\ninteraction_2 = client.interactions.create(\n agent=\"antigravity-preview-05-2026\",\n input=\"Modify the script to accept a name argument and greet the user.\",\n environment=interaction.environment_id,\n previous_interaction_id=interaction.id,\n)\nprint(interaction_2.output_text)\n" + }, + { + "label": "with_sources", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\ninteraction = client.interactions.create(\n agent=\"antigravity-preview-05-2026\",\n input=\"List all files under /workspace and summarize what you find.\",\n environment={\n \"type\": \"remote\",\n \"sources\": [\n {\n \"type\": \"repository\",\n \"source\": \"https://github.com/octocat/Spoon-Knife\",\n \"target\": \"/workspace/repo\",\n },\n {\n \"type\": \"inline\",\n \"content\": \"Focus on Python files only.\",\n \"target\": \"/workspace/notes.txt\",\n },\n ],\n },\n)\nprint(interaction.output_text)\n" + }, + { + "label": "custom_agent", + "lang": "python", + "source": "import uuid\nfrom google import genai\n\nclient = genai.Client()\n\n# Step 1: Create a custom agent.\nagent_id = f\"code-reviewer-{uuid.uuid4().hex[:8]}\"\nclient.agents.create(\n id=agent_id,\n base_agent=\"antigravity-preview-05-2026\",\n system_instruction=(\n \"You are a senior code reviewer. Check every file for bugs, style\"\n \" issues, and security vulnerabilities.\"\n ),\n base_environment={\n \"type\": \"remote\",\n \"sources\": [{\n \"type\": \"repository\",\n \"source\": \"https://github.com/octocat/Spoon-Knife\",\n \"target\": \"/workspace/repo\",\n }],\n },\n)\n\n# Step 2: Use the custom agent.\nresult = client.interactions.create(\n agent=agent_id,\n input=(\n \"Review the latest changes in /workspace/repo/src and file a\"\n \" summary.\"\n ),\n environment=\"remote\",\n)\nprint(result.output_text)\n" + }, + { + "label": "simple", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n input: 'Hello, how are you?',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "multi_turn", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n input: [\n {type: 'user_input', content: [{type: 'text', text: 'Hello'}]},\n {\n type: 'model_output',\n content: [\n {type: 'text', text: 'Hi there! How can I help you today?'},\n ],\n },\n {\n type: 'user_input',\n content: [{type: 'text', text: 'What is the capital of France?'}],\n },\n ],\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "multimodal_image", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n input: [\n {type: 'text', text: 'What is in this picture?'},\n {type: 'image', data: 'BASE64_ENCODED_IMAGE', mime_type: 'image/png'},\n ],\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "function_calling", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n tools: [\n {\n type: 'function',\n name: 'get_weather',\n description: 'Get the current weather in a given location',\n parameters: {\n type: 'object',\n properties: {\n location: {\n type: 'string',\n description: 'The city and state, e.g. San Francisco, CA',\n },\n },\n required: ['location'],\n },\n },\n ],\n input: 'What is the weather like in Boston, MA?',\n});\nconsole.log(interaction.steps.at(-1));\n" + }, + { + "label": "deep_research", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n agent: 'deep-research-pro-preview-12-2025',\n input: 'find a cure to cancer',\n background: true,\n});\nconsole.log(interaction.status);\n" + }, + { + "label": "antigravity", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n agent: 'antigravity-preview-05-2026',\n input:\n 'Read Hacker News, summarize the top 5 stories, and save results as a markdown file.',\n environment: 'remote',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "reuse_env", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n// Step 1: Create an interaction with a fresh remote environment.\nconst interaction = await ai.interactions.create({\n agent: 'antigravity-preview-05-2026',\n input: 'Write a hello world script at /workspace/hello.py.',\n environment: 'remote',\n});\nconsole.log(`Environment ID: ${interaction.environment_id}`);\n\n// Step 2: Reuse the same environment in a follow-up interaction.\nconst interaction2 = await ai.interactions.create({\n agent: 'antigravity-preview-05-2026',\n input: 'Modify the script to accept a name argument and greet the user.',\n environment: interaction.environment_id,\n previous_interaction_id: interaction.id,\n});\nconsole.log(interaction2.output_text);\n" + }, + { + "label": "with_sources", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n agent: 'antigravity-preview-05-2026',\n input: 'List all files under /workspace and summarize what you find.',\n environment: {\n type: 'remote',\n sources: [\n {\n type: 'repository',\n source: 'https://github.com/octocat/Spoon-Knife',\n target: '/workspace/repo',\n },\n {\n type: 'inline',\n content: 'Focus on Python files only.',\n target: '/workspace/notes.txt',\n },\n ],\n },\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "custom_agent", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n// Step 1: Create a custom agent.\nconst agentId = `code-reviewer-${crypto.randomUUID().slice(0, 8)}`;\nawait ai.agents.create({\n id: agentId,\n base_agent: 'antigravity-preview-05-2026',\n system_instruction:\n 'You are a senior code reviewer. Check every file for bugs, style issues, and security vulnerabilities.',\n base_environment: {\n type: 'remote',\n sources: [\n {\n type: 'repository',\n source: 'https://github.com/octocat/Spoon-Knife',\n target: '/workspace/repo',\n },\n ],\n },\n});\n\n// Step 2: Use the custom agent.\nconst result = await ai.interactions.create({\n agent: agentId,\n input:\n 'Review the latest changes in /workspace/repo/src and file a summary.',\n environment: 'remote',\n});\nconsole.log(result.output_text);\n" + }, + { + "label": "simple", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .input(InteractionsInput.of(\"Hello, how are you?\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "multi_turn", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.ModelOutputStep;\nimport com.google.genai.gaos.models.interactions.Step;\nimport com.google.genai.gaos.models.interactions.TextContent;\nimport com.google.genai.gaos.models.interactions.UserInputStep;\n\nClient client = new Client();\nList conversation =\n List.of(\n UserInputStep.builder()\n .content(List.of(TextContent.builder().text(\"Hello!\").build()))\n .build(),\n ModelOutputStep.builder()\n .content(\n List.of(\n TextContent.builder().text(\"Hi there! How can I help you today?\").build()))\n .build(),\n UserInputStep.builder()\n .content(\n List.of(TextContent.builder().text(\"What is the capital of France?\").build()))\n .build());\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .input(InteractionsInput.ofStep(conversation))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "multimodal_image", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.Content;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.ImageContent;\nimport com.google.genai.gaos.models.interactions.ImageContentMimeType;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.TextContent;\n\nClient client = new Client();\nList content =\n List.of(\n TextContent.builder().text(\"What is in this picture?\").build(),\n ImageContent.builder()\n .data(\"BASE64_ENCODED_IMAGE\")\n .mimeType(ImageContentMimeType.IMAGE_PNG)\n .build());\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .input(InteractionsInput.ofContent(content))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "function_calling", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Function;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.Step;\n\nClient client = new Client();\nMap parameters =\n Map.of(\n \"type\", \"object\",\n \"properties\",\n Map.of(\n \"location\",\n Map.of(\n \"type\", \"string\",\n \"description\", \"The city and state, e.g. San Francisco, CA\")),\n \"required\", List.of(\"location\"));\nFunction functionTool =\n Function.builder()\n .name(\"get_weather\")\n .description(\"Get the current weather in a given location\")\n .parameters(parameters)\n .build();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .tools(List.of(functionTool))\n .input(InteractionsInput.of(\"What is the weather like in Boston, MA?\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nList steps = interaction.steps().orElse(List.of());\nif (!steps.isEmpty()) {\n System.out.println(steps.get(steps.size() - 1));\n}\n" + }, + { + "label": "deep_research", + "lang": "java", + "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionStatus;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"deep-research-pro-preview-12-2025\"))\n .input(InteractionsInput.of(\"find a cure to cancer\"))\n .background(true)\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.status().map(InteractionStatus::value).orElse(\"\"));\n" + }, + { + "label": "antigravity", + "lang": "java", + "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteractionEnvironment;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"antigravity-preview-05-2026\"))\n .input(\n InteractionsInput.of(\n \"Read Hacker News, summarize the top 5 stories, and save results as a markdown\"\n + \" file.\"))\n .environment(CreateAgentInteractionEnvironment.of(\"remote\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "reuse_env", + "lang": "java", + "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteractionEnvironment;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.environments.Environment;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.Step;\n\nClient client = new Client();\n\n// Step 1: Create an interaction with a fresh remote environment.\nCreateAgentInteraction params1 =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"antigravity-preview-05-2026\"))\n .input(InteractionsInput.of(\"Write a hello world script at /workspace/hello.py.\"))\n .environment(CreateAgentInteractionEnvironment.of(\"remote\"))\n .build();\nCreateInteractionResponse response1 =\n client.interactions.create(CreateInteractionRequestBody.of(params1));\nInteraction interaction1 =\n response1.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(\"Environment ID: \" + interaction1.environmentId().orElse(\"\"));\n\n// Step 2: Reuse the same environment in a follow-up interaction.\nCreateAgentInteraction params2 =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"antigravity-preview-05-2026\"))\n .input(\n InteractionsInput.of(\n \"Modify the script to accept a name argument and greet the user.\"))\n .environment(\n CreateAgentInteractionEnvironment.of(interaction1.environmentId().orElse(\"\")))\n .previousInteractionId(interaction1.id().orElse(null))\n .build();\nCreateInteractionResponse response2 =\n client.interactions.create(CreateInteractionRequestBody.of(params2));\nInteraction interaction2 =\n response2.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction2.outputText().orElse(\"\"));\n" + }, + { + "label": "with_sources", + "lang": "java", + "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteractionEnvironment;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.environments.Environment;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.Source;\nimport com.google.genai.gaos.models.interactions.SourceType;\n\nClient client = new Client();\ncom.google.genai.gaos.models.interactions.Environment env =\n com.google.genai.gaos.models.interactions.Environment.builder()\n .sources(\n List.of(\n Source.builder()\n .type(SourceType.REPOSITORY)\n .source(\"https://github.com/octocat/Spoon-Knife\")\n .target(\"/workspace/repo\")\n .build(),\n Source.builder()\n .type(SourceType.INLINE)\n .content(\"Focus on Python files only.\")\n .target(\"/workspace/notes.txt\")\n .build()))\n .build();\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"antigravity-preview-05-2026\"))\n .input(\n InteractionsInput.of(\n \"List all files under /workspace and summarize what you find.\"))\n .environment(CreateAgentInteractionEnvironment.of(env))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "custom_agent", + "lang": "java", + "source": "import com.google.genai.gaos.models.agents.Agent;\nimport com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.gaos.models.agents.BaseEnvironment;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteractionEnvironment;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.environments.Environment;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.Source;\nimport com.google.genai.gaos.models.interactions.SourceType;\nimport com.google.genai.gaos.models.interactions.Step;\n\nClient client = new Client();\n\n// Step 1: Create a custom agent.\nString agentId = \"code-reviewer-\" + UUID.randomUUID().toString().substring(0, 8);\ncom.google.genai.gaos.models.interactions.Environment baseEnv =\n com.google.genai.gaos.models.interactions.Environment.builder()\n .sources(\n List.of(\n Source.builder()\n .type(SourceType.REPOSITORY)\n .source(\"https://github.com/octocat/Spoon-Knife\")\n .target(\"/workspace/repo\")\n .build()))\n .build();\nAgent customAgent =\n Agent.builder()\n .id(agentId)\n .baseAgent(\"antigravity-preview-05-2026\")\n .systemInstruction(\n \"You are a senior code reviewer. Check every file for bugs, style issues, and\"\n + \" security vulnerabilities.\")\n .baseEnvironment(BaseEnvironment.of(baseEnv))\n .build();\nclient.agents.create(customAgent);\n\n// Step 2: Use the custom agent.\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(agentId))\n .input(\n InteractionsInput.of(\n \"Review the latest changes in /workspace/repo/src and file a summary.\"))\n .environment(CreateAgentInteractionEnvironment.of(\"remote\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "simple", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Hello, how are you?\")),\n })\n\n res, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + }, + { + "label": "multi_turn", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n history := []gaos_interactions.Step{\n gaos_interactions.NewStep(gaos_interactions.UserInputStep{\n Content: []gaos_interactions.Content{\n gaos_interactions.NewContent(gaos_interactions.TextContent{\n Text: \"Hello!\",\n }),\n },\n }),\n gaos_interactions.NewStep(gaos_interactions.ModelOutputStep{\n Content: []gaos_interactions.Content{\n gaos_interactions.NewContent(gaos_interactions.TextContent{\n Text: \"Hi there! How can I help you today?\",\n }),\n },\n }),\n gaos_interactions.NewStep(gaos_interactions.UserInputStep{\n Content: []gaos_interactions.Content{\n gaos_interactions.NewContent(gaos_interactions.TextContent{\n Text: \"What is the capital of France?\",\n }),\n },\n }),\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(history)),\n })\n\n res, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + }, + { + "label": "multimodal_image", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n contents := []gaos_interactions.Content{\n gaos_interactions.NewContent(gaos_interactions.TextContent{\n Text: \"What is in this picture?\",\n }),\n gaos_interactions.NewContent(gaos_interactions.ImageContent{\n Data: genai.Ptr(\"BASE64_ENCODED_IMAGE\"),\n MimeType: gaos_interactions.ImageContentMimeTypeImagePng.ToPointer(),\n }),\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(contents)),\n })\n\n res, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + }, + { + "label": "function_calling", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n function := gaos_interactions.Function{\n Name: genai.Ptr(\"get_weather\"),\n Description: genai.Ptr(\"Get the current weather in a given location\"),\n Parameters: map[string]any{\n \"type\": \"object\",\n \"properties\": map[string]any{\n \"location\": map[string]any{\n \"type\": \"string\",\n \"description\": \"The city and state, e.g. San Francisco, CA\",\n },\n },\n \"required\": []string{\"location\"},\n },\n }\n\n tool := gaos_interactions.NewTool(function)\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"What is the weather like in Boston, MA?\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n res, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil {\n for _, step := range res.Interaction.Steps {\n if step.FunctionCallStep != nil {\n fmt.Println(\"Function Call:\", step.FunctionCallStep.Name)\n fmt.Println(\"Arguments:\", step.FunctionCallStep.Arguments)\n }\n }\n }\n}\n" + }, + { + "label": "deep_research", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"deep-research-pro-preview-12-2025\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"find a cure to cancer\")),\n Background: genai.Ptr(true),\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil {\n fmt.Println(res.Interaction.Status)\n }\n}\n" + }, + { + "label": "antigravity", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"antigravity-preview-05-2026\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Read Hacker News, summarize the top 5 stories, and save results as a markdown file.\")),\n Environment: genai.Ptr(gaos_interactions.NewCreateAgentInteractionEnvironment(\"remote\")),\n })\n\n res, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + }, + { + "label": "reuse_env", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n // Step 1: Create an interaction with a fresh remote environment.\n body1 := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"antigravity-preview-05-2026\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Write a hello world script at /workspace/hello.py.\")),\n Environment: genai.Ptr(gaos_interactions.NewCreateAgentInteractionEnvironment(\"remote\")),\n })\n\n res1, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body1})\n if err != nil {\n log.Fatal(err)\n }\n\n if res1.Interaction == nil || res1.Interaction.EnvironmentID == nil {\n log.Fatal(\"No environment ID returned\")\n }\n fmt.Printf(\"Environment ID: %s\\n\", *res1.Interaction.EnvironmentID)\n\n // Step 2: Reuse the same environment in a follow-up interaction.\n body2 := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"antigravity-preview-05-2026\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Modify the script to accept a name argument and greet the user.\")),\n Environment: genai.Ptr(gaos_interactions.NewCreateAgentInteractionEnvironment(*res1.Interaction.EnvironmentID)),\n PreviousInteractionID: res1.Interaction.ID,\n })\n\n res2, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body2})\n if err != nil {\n log.Fatal(err)\n }\n\n if res2.Interaction != nil && res2.Interaction.OutputText != nil {\n fmt.Println(*res2.Interaction.OutputText)\n }\n}\n" + }, + { + "label": "with_sources", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n env := gaos_interactions.Environment{\n Sources: []gaos_interactions.Source{\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeRepository),\n Source: genai.Ptr(\"https://github.com/octocat/Spoon-Knife\"),\n Target: genai.Ptr(\"/workspace/repo\"),\n },\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeInline),\n Content: genai.Ptr(\"Focus on Python files only.\"),\n Target: genai.Ptr(\"/workspace/notes.txt\"),\n },\n },\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"antigravity-preview-05-2026\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"List all files under /workspace and summarize what you find.\")),\n Environment: genai.Ptr(gaos_interactions.NewCreateAgentInteractionEnvironment(env)),\n })\n\n res, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + }, + { + "label": "custom_agent", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"time\"\n \"google.golang.org/genai\"\n gaos_agents \"google.golang.org/genai/interactions/models/agents\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n // Step 1: Create a custom agent.\n agentID := fmt.Sprintf(\"code-reviewer-%d\", time.Now().UnixNano())\n baseEnv := gaos_interactions.Environment{\n Sources: []gaos_interactions.Source{\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeRepository),\n Source: genai.Ptr(\"https://github.com/octocat/Spoon-Knife\"),\n Target: genai.Ptr(\"/workspace/repo\"),\n },\n },\n }\n\n _, err = client.Agents.Create(ctx, operations.CreateAgentRequest{\n Body: gaos_agents.Agent{\n ID: genai.Ptr(agentID),\n BaseAgent: genai.Ptr(\"antigravity-preview-05-2026\"),\n SystemInstruction: genai.Ptr(\"You are a senior code reviewer. Check every file for bugs, style issues, and security vulnerabilities.\"),\n BaseEnvironment: genai.Ptr(gaos_agents.NewBaseEnvironment(baseEnv)),\n },\n })\n if err != nil {\n log.Fatal(err)\n }\n\n // Step 2: Use the custom agent.\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(agentID),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Review the latest changes in /workspace/repo/src and file a summary.\")),\n Environment: genai.Ptr(gaos_interactions.NewCreateAgentInteractionEnvironment(\"remote\")),\n })\n\n res, err := client.Interactions.Create(ctx, operations.CreateInteractionRequest{Body: body})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + } + ] + } + }, + "/{api_version}/interactions/{id}": { + "delete": { + "operationId": "deleteInteraction", + "description": "Deletes the interaction by id.", + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + }, + { + "name": "id", + "in": "path", + "description": "The unique identifier of the interaction to delete.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "interactionsId", + "in": "path", + "description": "Required. The name of the interaction to delete.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Successful deletion of the interaction." + }, + "4XX": { + "description": "Error deleting interaction", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "$ref": "#/components/schemas/Error" + } + }, + "required": [ + "error" + ] + } + } + } + }, + "5XX": { + "description": "Error deleting interaction", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "$ref": "#/components/schemas/Error" + } + }, + "required": [ + "error" + ] + } + } + } + } + }, + "summary": "Deleting an interaction", + "x-speakeasy-exports": [ + { + "name": "InteractionDeleteParams", + "group": "interactions", + "representation": "input" + } + ], + "x-speakeasy-group": "interactions", + "x-speakeasy-name-override": "delete", + "x-codeSamples": [ + { + "label": "delete", + "lang": "sh", + "source": "curl -X DELETE \"https://generativelanguage.googleapis.com/v1beta/interactions/$INTERACTION_ID\" \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "delete", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n\nclient.interactions.delete(id=created.id)\nprint(\"Interaction deleted successfully.\")\n" + }, + { + "label": "delete", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nawait ai.interactions.delete(created.id);\nconsole.log('Interaction deleted successfully.');\n" + }, + { + "label": "delete", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.Interaction;\n\nClient client = new Client();\n\n\nclient.interactions.delete(interactionId);\nSystem.out.println(\"Interaction deleted successfully.\");\n" + }, + { + "label": "delete", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n req := operations.DeleteInteractionRequest{\n ID: interactionID,\n }\n\n _, err = client.Interactions.Delete(\n ctx,\n req,\n )\n if err != nil {\n log.Fatal(err)\n }\n\n fmt.Println(\"Interaction deleted successfully.\")\n}\n" + } + ] + }, + "get": { + "operationId": "getInteractionById", + "description": "Retrieves the full details of a single interaction based on its `Interaction.id`.", + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + }, + { + "name": "id", + "in": "path", + "description": "The unique identifier of the interaction to retrieve.", + "required": true, + "schema": { + "type": "string" + } + }, + { + "name": "last_event_id", + "in": "query", + "description": "Optional. If set, resumes the interaction stream from the next chunk after the event marked by the event id. Can only be used if `stream` is true.", + "schema": { + "type": "string" + } + }, + { + "name": "stream", + "in": "query", + "description": "If set to true, the generated content will be streamed incrementally.", + "schema": { + "type": "boolean", + "default": false + } + } + ], + "responses": { + "200": { + "description": "Successful retrieval of the interaction.", + "content": { + "application/json": { + "examples": { + "get": { + "summary": "Get Interaction", + "value": { + "created": "2025-11-26T12:25:15Z", + "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", + "model": "gemini-3.6-flash", + "object": "interaction", + "status": "completed", + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "I'm doing great, thank you for asking! How can I help you today?" + } + ] + } + ], + "updated": "2025-11-26T12:25:15Z" + } + } + }, + "schema": { + "$ref": "#/components/schemas/Interaction" + } + }, + "text/event-stream": { + "schema": { + "$ref": "#/components/schemas/InteractionSseStreamEvent" + }, + "x-speakeasy-sse-sentinel": "[DONE]" + } + } + }, + "4XX": { + "description": "Error getting interaction", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "$ref": "#/components/schemas/Error" + } + }, + "required": [ + "error" + ] + } + } + } + }, + "5XX": { + "description": "Error getting interaction", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "$ref": "#/components/schemas/Error" + } + }, + "required": [ + "error" + ] + } + } + } + } + }, + "summary": "Retrieving an interaction", + "x-speakeasy-exports": [ + { + "name": "InteractionGetParams", + "group": "interactions", + "representation": "input" + } + ], + "x-speakeasy-group": "interactions", + "x-speakeasy-name-override": "get", + "x-speakeasy-sse-overload": true, + "x-codeSamples": [ + { + "label": "get", + "lang": "sh", + "source": "curl -X GET \"https://generativelanguage.googleapis.com/v1beta/interactions/$INTERACTION_ID\" \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Api-Revision: 2026-05-20\"\n" + }, + { + "label": "get", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n\ninteraction = client.interactions.get(id=created.id)\nprint(interaction.status)\n" + }, + { + "label": "get", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n\nconst interaction = await ai.interactions.get(created.id);\nconsole.log(interaction.status);\n" + }, + { + "label": "get", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.GetInteractionByIdRequest;\nimport com.google.genai.gaos.models.operations.GetInteractionByIdResponse;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionStatus;\n\nClient client = new Client();\n\n\nGetInteractionByIdResponse getResponse =\n client.interactions.get(new GetInteractionByIdRequest(interactionId));\nInteraction interaction =\n getResponse\n .interaction()\n .orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.status().map(InteractionStatus::value).orElse(\"\"));\n" + }, + { + "label": "get", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n interaction, err := client.Interactions.Get(ctx, operations.GetInteractionByIDRequest{\n ID: *created.Interaction.ID,\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if interaction.Interaction != nil {\n fmt.Println(interaction.Interaction.Status)\n }\n}\n" + } + ] + } + }, + "/{api_version}/interactions/{id}/cancel": { + "post": { + "operationId": "cancelInteractionById", + "description": "Cancels an interaction by id. This only applies to background interactions that are still running.", + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + }, + { + "name": "id", + "in": "path", + "description": "The unique identifier of the interaction to cancel.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Successful operation", + "content": { + "application/json": { + "examples": { + "cancel": { + "summary": "Cancel Interaction", + "value": { + "created": "2025-11-26T12:25:15Z", + "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", + "model": "gemini-3.6-flash", + "object": "interaction", + "status": "cancelled", + "updated": "2025-11-26T12:25:15Z" + } + } + }, + "schema": { + "$ref": "#/components/schemas/Interaction" + } + } + } + }, + "4XX": { + "description": "Error cancelling interaction", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "$ref": "#/components/schemas/Error" + } + }, + "required": [ + "error" + ] + } + } + } + }, + "5XX": { + "description": "Error cancelling interaction", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "$ref": "#/components/schemas/Error" + } + }, + "required": [ + "error" + ] + } + } + } + } + }, + "summary": "Canceling an interaction", + "x-speakeasy-group": "interactions", + "x-speakeasy-name-override": "cancel", + "x-codeSamples": [ + { + "label": "cancel", + "lang": "sh", + "source": "curl -X POST \"https://generativelanguage.googleapis.com/v1beta/interactions/$INTERACTION_ID/cancel\" \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "cancel", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n# Start a background interaction so it stays in-progress.\ncreated = client.interactions.create(\n model=\"gemini-flash-latest\",\n input=\"Write a long essay about the history of computing.\",\n tools=[{\"type\": \"computer_use\"}],\n background=True,\n)\n\n# Cancel the in-progress interaction.\ninteraction = client.interactions.cancel(id=created.id)\nprint(interaction.status)\n" + }, + { + "label": "cancel", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n// Start a background interaction so it stays in-progress.\nconst created = await ai.interactions.create({\n model: 'gemini-flash-latest',\n input: 'Write a long essay about the history of computing.',\n tools: [{type: 'computer_use'}],\n background: true,\n});\n\n// Cancel the in-progress interaction.\nconst interaction = await ai.interactions.cancel(created.id);\nconsole.log(interaction.status);\n" + }, + { + "label": "cancel", + "lang": "java", + "source": "import com.google.genai.gaos.models.operations.CancelInteractionByIdResponse;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.ComputerUse;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionStatus;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\n\n// Start a background interaction so it stays in-progress.\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .input(InteractionsInput.of(\"Write a long essay about the history of computing.\"))\n .tools(List.of(new ComputerUse()))\n .background(true)\n .build();\nCreateInteractionResponse created =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nString interactionId = created.interaction().flatMap(Interaction::id).orElseThrow();\n\n// Cancel the in-progress interaction.\nCancelInteractionByIdResponse cancelResponse = client.interactions.cancel(interactionId);\nInteraction interaction =\n cancelResponse\n .interaction()\n .orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.status().map(InteractionStatus::value).orElse(\"\"));\n" + }, + { + "label": "cancel", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n req := operations.CancelInteractionByIDRequest{\n ID: interactionID,\n }\n\n res, err := client.Interactions.Cancel(\n ctx,\n req,\n )\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil {\n fmt.Println(res.Interaction.Status)\n }\n}\n" + } + ] + } + }, + "/{api_version}/triggers": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "get": { + "operationId": "ListTriggers", + "description": "Lists triggers for a project.", + "parameters": [ + { + "name": "filter", + "in": "query", + "description": "Optional. Filter expression (e.g., by state).", + "schema": { + "type": "string" + } + }, + { + "name": "page_size", + "in": "query", + "description": "Optional. The maximum number of triggers to return per page.", + "schema": { + "type": "integer", + "format": "int32" + } + }, + { + "name": "page_token", + "in": "query", + "description": "Optional. A page token from a previous ListTriggers call.", + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListTriggersResponse" + } + } + } + } + }, + "tags": [ + "triggers" + ] + }, + "post": { + "operationId": "CreateTrigger", + "description": "Creates a new trigger that will invoke the specified agent on the given\ncron schedule.", + "parameters": [], + "requestBody": { + "description": "Required. The trigger configuration to create.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Trigger" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Trigger" + } + } + } + } + }, + "tags": [ + "triggers" + ], + "x-codeSamples": [ + { + "label": "create", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/triggers \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"schedule\": \"0 9 * * *\",\n \"time_zone\": \"America/New_York\",\n \"interaction\": {\n \"agent\": \"antigravity-preview-05-2026\",\n \"input\": \"Summarize top news stories.\",\n \"environment\": \"remote\"\n }\n }'\n" + }, + { + "label": "create", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\ntrigger = client.triggers.create(\n schedule=\"0 9 * * *\",\n time_zone=\"America/New_York\",\n interaction={\n \"agent\": \"antigravity-preview-05-2026\",\n \"input\": \"Summarize top news stories.\",\n \"environment\": \"remote\",\n },\n)\nprint(trigger.id, trigger.interaction.agent)\n" + }, + { + "label": "create", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst trigger = await ai.triggers.create({\n schedule: '0 9 * * *',\n time_zone: 'America/New_York',\n interaction: {\n agent: 'antigravity-preview-05-2026',\n input: 'Summarize top news stories.',\n environment: 'remote',\n },\n});\nconsole.log(trigger.id, trigger.interaction.agent);\n" + }, + { + "label": "create", + "lang": "java", + "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteractionEnvironment;\nimport com.google.genai.gaos.models.operations.CreateTriggerResponse;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.triggers.Trigger;\nimport com.google.genai.gaos.models.triggers.TriggerCreateParams;\n\nClient client = new Client();\nCreateAgentInteraction interaction =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"antigravity-preview-05-2026\"))\n .input(InteractionsInput.of(\"Summarize top news stories.\"))\n .environment(CreateAgentInteractionEnvironment.of(\"remote\"))\n .build();\nTriggerCreateParams params =\n TriggerCreateParams.builder()\n .schedule(\"0 9 * * *\")\n .timeZone(\"America/New_York\")\n .interaction(interaction)\n .build();\nCreateTriggerResponse response = client.triggers.create(params);\nTrigger trigger =\n response.trigger().orElseThrow(() -> new RuntimeException(\"No trigger returned\"));\nSystem.out.println(\n trigger.id().orElse(\"\")\n + \" \"\n + trigger.interaction().flatMap(CreateAgentInteraction::agent).orElse(null));\n" + } + ] + } + }, + "/{api_version}/triggers/{id}": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "delete": { + "operationId": "DeleteTrigger", + "description": "Deletes a trigger.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. Resource name of the trigger.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Empty" + } + } + } + } + }, + "tags": [ + "triggers" + ] + }, + "get": { + "operationId": "GetTrigger", + "description": "Gets details of a single trigger.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. Resource name of the trigger.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Trigger" + } + } + } + } + }, + "tags": [ + "triggers" + ] + }, + "patch": { + "operationId": "UpdateTrigger", + "description": "Updates a trigger.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. Resource name of the trigger.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "description": "Required. The trigger with updated fields.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/UpdateTriggerRequest" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Trigger" + } + } + } + } + }, + "tags": [ + "triggers" + ] + } + }, + "/{api_version}/triggers/{triggerId}/executions": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "get": { + "operationId": "ListTriggerExecutions", + "description": "Lists executions for a trigger.", + "parameters": [ + { + "name": "filter", + "in": "query", + "description": "Optional. Filter expression.", + "schema": { + "type": "string" + } + }, + { + "name": "page_size", + "in": "query", + "description": "The maximum number of executions to return per page.", + "schema": { + "type": "integer", + "format": "int32" + } + }, + { + "name": "page_token", + "in": "query", + "description": "A page token from a previous ListTriggerExecutions call.", + "schema": { + "type": "string" + } + }, + { + "name": "triggerId", + "in": "path", + "description": "Required. The trigger ID to list executions from.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListTriggerExecutionsResponse" + } + } + } + } + }, + "tags": [ + "triggers.executions" + ] + }, + "post": { + "operationId": "RunTrigger", + "description": "Runs a trigger immediately.", + "parameters": [ + { + "name": "triggerId", + "in": "path", + "description": "Required. Resource name of the trigger.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TriggerExecution" + } + } + } + } + }, + "tags": [ + "triggers" + ] + } + }, + "/{api_version}/voices": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "get": { + "operationId": "ListVoices", + "description": "Lists custom stored voices owned by the caller (ordered newest first)\nfollowed by prebuilt system voices from Google's voice catalog.", + "parameters": [ + { + "name": "accent", + "in": "query", + "description": "Optional. Filter by accent description (e.g. \"American\", \"British\").\nCase-insensitive exact match. If multiple values are specified, matches\nvoices with any of the specified accents (OR).", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "name": "context", + "in": "query", + "description": "Optional. Filter by intended context or domain (e.g. \"News, Commercial\").\nCase-insensitive exact match. If multiple values are specified, matches\nvoices with any of the specified contexts (OR).", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "name": "gender", + "in": "query", + "description": "Optional. Filter by gender presentation (e.g. \"female\", \"male\", \"neutral\").\nCase-insensitive exact match. If multiple values are specified, matches\nvoices with any of the specified genders (OR).", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "name": "language_code", + "in": "query", + "description": "Optional. Filter by BCP-47 language code (e.g. \"en-US\").\nCase-insensitive exact match. If multiple values are specified, matches\nvoices with any of the specified language codes (OR).", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "name": "page_size", + "in": "query", + "description": "Optional. The maximum number of voices to return per page. The service may\nreturn fewer than this value. If unspecified, at most 50 voices are\nreturned. The maximum value is 1000; values above 1000 are coerced to\n1000.", + "schema": { + "type": "integer", + "format": "int32" + } + }, + { + "name": "page_token", + "in": "query", + "description": "Optional. A page token received from a previous `ListVoices` call. Provide\nthis to retrieve the subsequent page.\n\nWhen paginating, all filter query parameters (`language_code`,\n`region_code`, `accent`, `persona`, `context`, `gender`, `pitch`, `type`,\nand `search`) must match the call that returned this token; otherwise the\nrequest fails with `INVALID_ARGUMENT`. `page_size` may change between\npages.", + "schema": { + "type": "string" + } + }, + { + "name": "persona", + "in": "query", + "description": "Optional. Filter by vocal persona (e.g. \"Warm, Friendly\").\nCase-insensitive exact match. If multiple values are specified, matches\nvoices with any of the specified personas (OR).", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "name": "pitch", + "in": "query", + "description": "Optional. Filter by vocal pitch. Accepts `\"low\"`, `\"medium\"`, `\"high\"`\n(case-insensitive). If multiple values are specified, matches voices with\nany of the specified pitches (OR).", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "name": "region_code", + "in": "query", + "description": "Optional. Filter by ISO 3166-1 alpha-2 or UN M.49 region code (e.g. \"US\",\n\"001\").\nCase-insensitive exact match. If multiple values are specified, matches\nvoices with any of the specified region codes (OR).", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "name": "search", + "in": "query", + "description": "Optional. Free-text substring search query matched case-insensitively\nagainst both `display_name` and `description`. Maximum 2048 bytes.", + "schema": { + "type": "string" + } + }, + { + "name": "type", + "in": "query", + "description": "Optional. Filter by voice type. Accepts `\"prebuilt\"`, `\"replicated\"`,\n`\"prompted\"` (case-insensitive). If multiple values are specified, matches\nvoices with any of the specified types (OR).", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListVoicesResponse" + } + } + } + } + }, + "tags": [ + "voices" + ], + "x-codeSamples": [ + { + "label": "list", + "lang": "sh", + "source": "curl -X GET https://generativelanguage.googleapis.com/v1beta/voices \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "list", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.voices.list()\nfor voice in response.voices or []:\n print(voice.id, voice.display_name)\n" + }, + { + "label": "list", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst response = await ai.voices.list();\nfor (const voice of response.voices ?? []) {\n console.log(voice.id, voice.display_name);\n}\n" + }, + { + "label": "list", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n res, err := client.Voices.List(ctx, operations.ListVoicesRequest{})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.ListVoicesResponse != nil {\n for _, voice := range res.ListVoicesResponse.Voices {\n if voice.ID != nil && voice.DisplayName != nil {\n fmt.Println(*voice.ID, *voice.DisplayName)\n }\n }\n }\n}\n" + } + ] + }, + "post": { + "operationId": "CreateVoice", + "description": "Creates a custom voice from a natural-language prompt\n(`VOICE_TYPE_PROMPTED`) or from reference and consent audio recordings\n(`VOICE_TYPE_REPLICATED`).", + "parameters": [], + "requestBody": { + "description": "The request body.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CreateVoiceRequest" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Voice" + } + } + } + } + }, + "tags": [ + "voices" + ], + "x-codeSamples": [ + { + "label": "create", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/voices \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"store\": true,\n \"voice\": {\n \"type\": \"prompted\",\n \"display_name\": \"Warm Narrator\",\n \"language_code\": \"en-US\",\n \"prompted\": {\n \"input\": \"A warm, friendly narrator voice with a calm pace.\"\n }\n }\n }'\n" + }, + { + "label": "create", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nvoice = client.voices.create(\n store=True,\n voice={\n \"type\": \"prompted\",\n \"display_name\": \"Warm Narrator\",\n \"language_code\": \"en-US\",\n \"prompted\": {\n \"input\": \"A warm, friendly narrator voice with a calm pace.\",\n },\n },\n)\nprint(voice.id)\n" + }, + { + "label": "create", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst voice = await ai.voices.create({\n store: true,\n voice: {\n type: 'prompted',\n display_name: 'Warm Narrator',\n language_code: 'en-US',\n prompted: {\n input: 'A warm, friendly narrator voice with a calm pace.',\n },\n },\n});\nconsole.log(voice.id);\n" + }, + { + "label": "create", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_voices \"google.golang.org/genai/interactions/models/voices\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n req := operations.CreateVoiceRequest{\n Body: gaos_voices.CreateVoiceRequest{\n Store: genai.Ptr(true),\n Voice: gaos_voices.VoiceInput{\n Type: gaos_voices.VoiceTypePrompted,\n DisplayName: genai.Ptr(\"Warm Narrator\"),\n LanguageCode: genai.Ptr(\"en-US\"),\n Prompted: &gaos_voices.PromptedVoice{\n Input: \"A warm, friendly narrator voice with a calm pace.\",\n },\n },\n },\n }\n\n res, err := client.Voices.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Voice != nil && res.Voice.ID != nil {\n fmt.Println(*res.Voice.ID)\n }\n}\n" + } + ] + } + }, + "/{api_version}/voices/{voicesId}": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "delete": { + "operationId": "DeleteVoice", + "description": "Deletes a custom stored voice (`store = true`) by resource name.\nPrebuilt catalog voices (`VOICE_TYPE_PREBUILT`) cannot be deleted.", + "parameters": [ + { + "name": "voicesId", + "in": "path", + "description": "Required. The resource name of the custom stored voice to delete\n(for example, `voices/voice_abc123def456`).", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DeleteVoiceResponse" + } + } + } + } + }, + "tags": [ + "voices" + ], + "x-codeSamples": [ + { + "label": "delete", + "lang": "sh", + "source": "curl -X DELETE https://generativelanguage.googleapis.com/v1beta/voices/voice_abc123 \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "delete", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nclient.voices.delete(id=\"voice_abc123\")\n" + }, + { + "label": "delete", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nawait ai.voices.delete('voice_abc123');\n" + }, + { + "label": "delete", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n _, err = client.Voices.Delete(ctx, operations.DeleteVoiceRequest{\n ID: \"voice_abc123\",\n })\n if err != nil {\n log.Fatal(err)\n }\n}\n" + } + ] + }, + "get": { + "operationId": "GetVoice", + "description": "Gets a custom stored voice (`store = true`) by resource name.\nPrebuilt catalog voices (`VOICE_TYPE_PREBUILT`) cannot be retrieved via\n`GetVoice`; use `ListVoices` instead.", + "parameters": [ + { + "name": "voicesId", + "in": "path", + "description": "Required. The resource name of the custom stored voice to retrieve\n(for example, `voices/voice_abc123def456`).", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Voice" + } + } + } + } + }, + "tags": [ + "voices" + ], + "x-codeSamples": [ + { + "label": "get", + "lang": "sh", + "source": "curl -X GET https://generativelanguage.googleapis.com/v1beta/voices/voice_abc123 \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\"\n" + }, + { + "label": "get", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nvoice = client.voices.get(id=\"voice_abc123\")\nprint(voice.id, voice.display_name)\n" + }, + { + "label": "get", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst voice = await ai.voices.get('voice_abc123');\nconsole.log(voice.id, voice.display_name);\n" + }, + { + "label": "get", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n res, err := client.Voices.Get(ctx, operations.GetVoiceRequest{\n ID: \"voice_abc123\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Voice != nil && res.Voice.ID != nil && res.Voice.DisplayName != nil {\n fmt.Println(*res.Voice.ID, *res.Voice.DisplayName)\n }\n}\n" + } + ] + } + }, + "/{api_version}/webhooks": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "get": { + "operationId": "ListWebhooks", + "description": "Lists all Webhooks.", + "parameters": [ + { + "name": "page_size", + "in": "query", + "description": "Optional. The maximum number of webhooks to return. The service may return\nfewer than this value. If unspecified, at most 50 webhooks will be\nreturned. The maximum value is 1000.", + "schema": { + "type": "integer", + "format": "int32" + } + }, + { + "name": "page_token", + "in": "query", + "description": "Optional. A page token, received from a previous `ListWebhooks` call.\nProvide this to retrieve the subsequent page.", + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListWebhooksResponse" + } + } + } + } + }, + "tags": [ + "webhooks" + ] + }, + "post": { + "operationId": "CreateWebhook", + "description": "Creates a new Webhook.", + "parameters": [ + { + "name": "webhook_id", + "in": "query", + "description": "The webhook_id to use for the webhook. If not specified, the server will\ngenerate a unique ID.", + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "description": "Required. The request body.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CreateWebhookRequest" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Webhook" + } + } + } + } + }, + "tags": [ + "webhooks" + ] + } + }, + "/{api_version}/webhooks/{id}": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "delete": { + "operationId": "DeleteWebhook", + "description": "Deletes a Webhook.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. The ID of the webhook to delete.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Empty" + } + } + } + } + }, + "tags": [ + "webhooks" + ] + }, + "get": { + "operationId": "GetWebhook", + "description": "Gets a specific Webhook.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. The ID of the webhook to retrieve.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Webhook" + } + } + } + } + }, + "tags": [ + "webhooks" + ] + }, + "patch": { + "operationId": "UpdateWebhook", + "description": "Updates an existing Webhook.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. The ID of the webhook to update.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "description": "Required. The request body.", + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/UpdateWebhookRequest" + } + } + } + }, + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Webhook" + } + } + } + } + }, + "tags": [ + "webhooks" + ] + } + }, + "/{api_version}/webhooks/{id}:ping": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "post": { + "operationId": "PingWebhook", + "description": "Sends a ping event to a Webhook.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. The ID of the webhook to ping.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/PingWebhookResponse" + } + } + } + } + }, + "tags": [ + "webhooks" + ] + } + }, + "/{api_version}/webhooks/{id}:rotateSigningSecret": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + } + ], + "post": { + "operationId": "RotateSigningSecret", + "description": "Generates a new signing secret for a Webhook.", + "parameters": [ + { + "name": "id", + "in": "path", + "description": "Required. The ID of the webhook for which to generate a signing secret.", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "default": { + "description": "Successful operation", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RotateSigningSecretResponse" + } + } + } + } + }, + "tags": [ + "webhooks" + ] + } + } + }, + "components": { + "parameters": { + "EnvironmentFileExtract": { + "name": "extract", + "in": "query", + "description": "Optional. If true, treats the uploaded file as a tar/tar.gz archive and unpacks it into `path`.", + "schema": { + "type": "boolean" + } + }, + "EnvironmentFileOverwrite": { + "name": "overwrite", + "in": "query", + "description": "Optional. Whether to overwrite the destination file if it already exists.", + "schema": { + "type": "boolean" + } + }, + "EnvironmentFileUploadCommand": { + "name": "X-Goog-Upload-Command", + "in": "header", + "description": "Command that starts the resumable upload session.", + "required": true, + "schema": { + "type": "string", + "const": "start", + "default": "start" + } + }, + "EnvironmentFileUploadContentLength": { + "name": "X-Goog-Upload-Header-Content-Length", + "in": "header", + "description": "Total number of file bytes that will be uploaded to the session URL.", + "required": true, + "schema": { + "type": "integer", + "format": "int64", + "minimum": 0 + } + }, + "EnvironmentFileUploadContentType": { + "name": "X-Goog-Upload-Header-Content-Type", + "in": "header", + "description": "MIME type of the file that will be uploaded to the session URL.", + "required": true, + "schema": { + "type": "string", + "minLength": 1 + } + }, + "EnvironmentFileUploadProtocol": { + "name": "X-Goog-Upload-Protocol", + "in": "header", + "description": "Resumable upload protocol selector.", + "required": true, + "schema": { + "type": "string", + "const": "resumable", + "default": "resumable" + } + }, + "api_version": { + "name": "api_version", + "in": "path", + "description": "API version for request routing.", + "required": true, + "schema": { + "type": "string" + } + }, + "google_genai_user_project": { + "name": "x-goog-user-project", + "in": "header", + "description": "Quota project header to send with Google GenAI API requests.", + "schema": { + "type": "string" + }, + "x-speakeasy-name-override": "user_project" + } + }, + "headers": { + "EnvironmentFileUploadUrl": { + "description": "Opaque URL to which file chunks are uploaded and finalized.", + "schema": { + "type": "string", + "format": "uri" + } + } + }, + "schemas": { + "Agent": { + "description": "An agent definition for the CreateAgent API.\nThis message is the target for annotation-parser-based JSON parsing.\nNew format:\n {\n \"id\": \"customer-sentinel\",\n \"base_agent\": \"\",\n \"system_instruction\": \"...\",\n \"base_environment\": { \"type\": \"remote\", \"sources\": [...] },\n \"tools\": [ {\"type\": \"code_execution\"} ]\n }", + "type": "object", + "properties": { + "agent_config": { + "description": "Configuration parameters for the agent.", + "oneOf": [ + { + "$ref": "#/components/schemas/AntigravityAgentConfig" + } + ] + }, + "base_environment": { + "description": "The environment configuration for the agent.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentConfig" + }, + { + "type": "string" + } + ] + }, + "id": { + "description": "The unique identifier for the agent.", + "type": "string" + }, + "base_agent": { + "description": "The base agent to extend.", + "type": "string" + }, + "system_instruction": { + "description": "System instruction for the agent.", + "type": "string" + }, + "description": { + "description": "Agent description for developers to quickly read and understand.", + "type": "string" + }, + "tools": { + "description": "The tools available to the agent.", + "type": "array", + "items": { + "$ref": "#/components/schemas/AgentTool" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "AgentOption": { + "title": "Agent", + "description": "The agent to interact with.", + "type": "string", + "enum": [ + "deep-research-pro-preview-12-2025", + "deep-research-preview-04-2026", + "deep-research-max-preview-04-2026", + "antigravity-preview-05-2026" + ], + "x-speakeasy-enum-descriptions": [ + "Gemini Deep Research Agent", + "Gemini Deep Research Agent", + "Gemini Deep Research Max Agent", + "Use the Antigravity managed agent to perform multi-step tasks that require reasoning, file operations, and tool use." + ], + "x-speakeasy-enum-format": "union", + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-unknown-values": "allow" + }, + "AgentTool": { + "description": "A tool that the agent can use.", + "oneOf": [ + { + "$ref": "#/components/schemas/CodeExecution" + }, + { + "$ref": "#/components/schemas/Function" + }, + { + "$ref": "#/components/schemas/GoogleSearch" + }, + { + "$ref": "#/components/schemas/McpServer" + }, + { + "$ref": "#/components/schemas/UrlContext" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "AllowedTools": { + "description": "The configuration for allowed tools.", + "type": "object", + "properties": { + "mode": { + "type": "string", + "enum": [ + "auto", + "any", + "none", + "validated" + ], + "x-google-enum-descriptions": [ + "Auto tool choice.", + "Any tool choice.", + "No tool choice.", + "Validated tool choice." + ], + "x-speakeasy-model-namespace": "interactions", + "description": "The mode of the tool choice." + }, + "tools": { + "description": "The names of the allowed tools.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "Annotation": { + "description": "Citation information for model-generated content.", + "type": "object", + "oneOf": [ + { + "$ref": "#/components/schemas/FileCitation" + }, + { + "$ref": "#/components/schemas/PlaceCitation" + }, + { + "$ref": "#/components/schemas/SpeechAnnotation" + }, + { + "$ref": "#/components/schemas/UrlCitation" + }, + { + "$ref": "#/components/schemas/WordInfo" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "AntigravityAgentConfig": { + "description": "Configuration for the Antigravity agent runtime.\nProvides server-side control over the agent's execution environment\nand tool configuration.", + "type": "object", + "properties": { + "type": { + "const": "antigravity" + }, + "model": { + "description": "The model to use for agent reasoning.", + "type": "string" + }, + "max_total_tokens": { + "description": "Max total tokens for the agent run.", + "type": "string", + "format": "int64" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ArgumentsDelta": { + "type": "object", + "properties": { + "type": { + "const": "arguments_delta" + }, + "arguments": { + "type": "string" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "AudioContent": { + "description": "An audio content block.", + "type": "object", + "properties": { + "type": { + "const": "audio" + }, + "data": { + "description": "The audio content.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "description": "The URI of the audio.", + "type": "string" + }, + "channels": { + "description": "The number of audio channels.", + "type": "integer", + "format": "int32" + }, + "sample_rate": { + "description": "The sample rate of the audio.", + "type": "integer", + "format": "int32" + }, + "mime_type": { + "description": "The mime type of the audio.", + "type": "string", + "enum": [ + "audio/wav", + "audio/mp3", + "audio/aiff", + "audio/aac", + "audio/ogg", + "audio/flac", + "audio/mpeg", + "audio/m4a", + "audio/l16", + "audio/opus", + "audio/alaw", + "audio/mulaw", + "audio/webm" + ], + "x-google-enum-descriptions": [ + "WAV audio format", + "MP3 audio format", + "AIFF audio format", + "AAC audio format", + "OGG audio format", + "FLAC audio format", + "MPEG audio format", + "M4A audio format", + "L16 audio format", + "OPUS audio format", + "ALAW audio format", + "MULAW audio format", + "WEBM audio format" + ] + } + }, + "required": [ + "type" + ], + "examples": [ + { + "audio": { + "summary": "Audio", + "value": { + "type": "audio", + "data": "BASE64_ENCODED_AUDIO", + "mime_type": "audio/wav" + } + } + } + ], + "x-speakeasy-exports": [ + { + "name": "audio_content_param", + "group": "interactions", + "representation": "input" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "AudioData": { + "description": "Audio payload used for voice creation.", + "type": "object", + "properties": { + "mime_type": { + "description": "Required. The IANA MIME type of the audio data (for example, `audio/wav` or\n`audio/mpeg`).", + "type": "string" + }, + "data": { + "description": "Required. The raw audio bytes.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "data", + "mime_type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "AudioDelta": { + "type": "object", + "properties": { + "type": { + "const": "audio" + }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, + "mime_type": { + "type": "string", + "enum": [ + "audio/wav", + "audio/mp3", + "audio/aiff", + "audio/aac", + "audio/ogg", + "audio/flac", + "audio/mpeg", + "audio/m4a", + "audio/l16", + "audio/opus", + "audio/alaw", + "audio/mulaw", + "audio/webm" + ], + "x-google-enum-descriptions": [ + "WAV audio format", + "MP3 audio format", + "AIFF audio format", + "AAC audio format", + "OGG audio format", + "FLAC audio format", + "MPEG audio format", + "M4A audio format", + "L16 audio format", + "OPUS audio format", + "ALAW audio format", + "MULAW audio format", + "WEBM audio format" + ] + }, + "rate": { + "description": "Deprecated. Use sample_rate instead. The value is ignored.", + "type": "integer", + "deprecated": true, + "format": "int32" + }, + "sample_rate": { + "description": "The sample rate of the audio.", + "type": "integer", + "format": "int32" + }, + "channels": { + "description": "The number of audio channels.", + "type": "integer", + "format": "int32" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "AudioResponseFormat": { + "description": "Configuration for audio output format.", + "type": "object", + "properties": { + "type": { + "const": "audio" + }, + "mime_type": { + "description": "The MIME type of the audio output.", + "type": "string", + "enum": [ + "audio/mp3", + "audio/ogg_opus", + "audio/l16", + "audio/wav", + "audio/alaw", + "audio/mulaw" + ], + "x-google-enum-descriptions": [ + "MP3 audio format.", + "OGG Opus audio format.", + "Raw PCM (L16) audio format.", + "WAV audio format.", + "A-law audio format.", + "Mu-law audio format." + ] + }, + "delivery": { + "description": "The delivery mode for the audio output.", + "type": "string", + "enum": [ + "inline", + "uri" + ], + "x-google-enum-descriptions": [ + "Audio data is returned inline in the response.", + "Audio data is returned as a URI." + ] + }, + "sample_rate": { + "description": "Sample rate in Hz.", + "type": "integer", + "format": "int32" + }, + "bit_rate": { + "description": "Bit rate in bits per second (bps). Only applicable for compressed formats\n(MP3, Opus).", + "type": "integer", + "format": "int32" + } + }, + "required": [ + "type" + ], + "examples": [ + { + "audio_response_format": { + "summary": "Audio Output", + "value": { + "type": "audio", + "sample_rate": 24000 + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "CodeExecution": { + "description": "A tool that can be used by the model to execute code.", + "type": "object", + "properties": { + "type": { + "const": "code_execution" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-codeSamples": [ + { + "label": "code_execution", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"tools\": [{\n \"type\": \"code_execution\"\n }],\n \"input\": \"Calculate the first 10 Fibonacci numbers\"\n }'\n" + }, + { + "label": "code_execution", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n tools=[{\"type\": \"code_execution\"}],\n input=\"Calculate the first 10 Fibonacci numbers\",\n)\nprint(response.output_text)\n" + }, + { + "label": "code_execution", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n tools: [{type: 'code_execution'}],\n input: 'Calculate the first 10 Fibonacci numbers',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "code_execution", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CodeExecution;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .tools(List.of(new CodeExecution()))\n .input(InteractionsInput.of(\"Calculate the first 10 Fibonacci numbers\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "code_execution", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n tool := gaos_interactions.NewTool(gaos_interactions.CodeExecution{})\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Calculate the first 10 Fibonacci numbers\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + } + ] + }, + "CodeExecutionCallArguments": { + "description": "The arguments to pass to the code execution.", + "type": "object", + "properties": { + "language": { + "description": "Programming language of the `code`.", + "type": "string", + "enum": [ + "python", + "bash" + ], + "x-google-enum-descriptions": [ + "Python >= 3.10, with numpy and simpy available.", + "Bash. Only available for Gemini 3.0 model or above." + ] + }, + "code": { + "description": "The code to be executed.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "CodeExecutionCallDelta": { + "type": "object", + "properties": { + "arguments": { + "$ref": "#/components/schemas/CodeExecutionCallArguments" + }, + "type": { + "const": "code_execution_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "arguments", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "CodeExecutionCallStep": { + "description": "Code execution call step.", + "type": "object", + "properties": { + "arguments": { + "description": "Required. The arguments to pass to the code execution.", + "$ref": "#/components/schemas/CodeExecutionCallStepArguments" + }, + "id": { + "description": "Required. A unique ID for this specific tool call.", + "type": "string" + }, + "type": { + "const": "code_execution_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "arguments", + "id", + "type" + ], + "examples": [ + { + "code_execution_call": { + "summary": "CodeExecutionCallStep", + "value": { + "type": "code_execution_call", + "arguments": { + "code": "print(sum(range(1, 11)))" + }, + "id": "code_call_71021" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "CodeExecutionCallStepArguments": { + "description": "The arguments to pass to the code execution.", + "type": "object", + "properties": { + "language": { + "description": "Programming language of the `code`.", + "type": "string", + "enum": [ + "python", + "bash" + ], + "x-google-enum-descriptions": [ + "Python >= 3.10, with numpy and simpy available.", + "Bash. Only available for Gemini 3.0 model or above." + ] + }, + "code": { + "description": "The code to be executed.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "CodeExecutionCallArguments" + }, + "CodeExecutionResultDelta": { + "type": "object", + "properties": { + "result": { + "type": "string" + }, + "type": { + "const": "code_execution_result" + }, + "is_error": { + "type": "boolean" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "result", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "CodeExecutionResultStep": { + "description": "Code execution result step.", + "type": "object", + "properties": { + "result": { + "description": "Required. The output of the code execution.", + "type": "string" + }, + "call_id": { + "description": "Required. ID to match the ID from the function call block.", + "type": "string" + }, + "type": { + "const": "code_execution_result" + }, + "is_error": { + "description": "Whether the code execution resulted in an error.", + "type": "boolean" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "call_id", + "result", + "type" + ], + "examples": [ + { + "code_execution_result": { + "summary": "CodeExecutionResultStep", + "value": { + "type": "code_execution_result", + "call_id": "code_call_71021", + "result": "55\n" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ComputerUse": { + "description": "A tool that can be used by the model to interact with the computer.", + "type": "object", + "properties": { + "type": { + "const": "computer_use" + }, + "environment": { + "description": "The environment being operated.", + "type": "string", + "enum": [ + "browser", + "mobile", + "desktop" + ], + "x-google-enum-descriptions": [ + "Operates in a web browser.", + "Operates in a mobile environment.", + "Operates in a desktop environment." + ] + }, + "excluded_predefined_functions": { + "description": "The list of predefined functions that are excluded from the model call.", + "type": "array", + "items": { + "type": "string" + } + }, + "enable_prompt_injection_detection": { + "description": "Whether enable the prompt injection detection check on computer-use\nrequest.", + "type": "boolean" + }, + "disabled_safety_policies": { + "description": "Optional. Disabled safety policies for computer use.", + "type": "array", + "items": { + "type": "string", + "enum": [ + "financial_transactions", + "sensitive_data_modification", + "communication_tool", + "account_creation", + "data_modification", + "user_consent_management", + "legal_terms_and_agreements" + ], + "x-google-enum-descriptions": [ + "Safety policy for financial transactions.", + "Safety policy for sensitive data modification.", + "Safety policy for communication tools (e.g. Gmail, Chat, Meet).", + "Safety policy for account creation.", + "Safety policy for data modification.", + "Safety policy for user consent management.", + "Safety policy for legal terms and agreements." + ] + } + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-codeSamples": [ + { + "label": "computer_use", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-2.5-computer-use-preview-10-2025\",\n \"tools\": [{\n \"type\": \"computer_use\"\n }],\n \"input\": \"Find a flight to Tokyo\"\n }'\n" + }, + { + "label": "computer_use", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-2.5-computer-use-preview-10-2025\",\n tools=[{\"type\": \"computer_use\"}],\n input=\"Find a flight to Tokyo\",\n)\nprint(response.output_text)\n" + }, + { + "label": "computer_use", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-2.5-computer-use-preview-10-2025',\n tools: [{type: 'computer_use'}],\n input: 'Find a flight to Tokyo',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "computer_use", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.ComputerUse;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-2.5-computer-use-preview-10-2025\")\n .tools(List.of(new ComputerUse()))\n .input(InteractionsInput.of(\"Find a flight to Tokyo\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "computer_use", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n tool := gaos_interactions.NewTool(gaos_interactions.ComputerUse{})\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-2.5-computer-use-preview-10-2025\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Find a flight to Tokyo\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + } + ] + }, + "Content": { + "description": "The content of the response.", + "type": "object", + "oneOf": [ + { + "$ref": "#/components/schemas/AudioContent" + }, + { + "$ref": "#/components/schemas/DocumentContent" + }, + { + "$ref": "#/components/schemas/ImageContent" + }, + { + "$ref": "#/components/schemas/TextContent" + }, + { + "$ref": "#/components/schemas/VideoContent" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "CreateAgentInteractionParams": { + "description": "Parameters for creating agent interactions", + "properties": { + "agent": { + "description": "The name of the `Agent` used for generating the interaction.", + "$ref": "#/components/schemas/AgentOption" + }, + "input": { + "$ref": "#/components/schemas/InteractionsInput" + }, + "agent_config": { + "description": "Configuration parameters for the agent interaction.", + "oneOf": [ + { + "$ref": "#/components/schemas/AntigravityAgentConfig" + }, + { + "$ref": "#/components/schemas/DeepResearchAgentConfig" + }, + { + "$ref": "#/components/schemas/DynamicAgentConfig" + } + ] + }, + "response_format": { + "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", + "oneOf": [ + { + "$ref": "#/components/schemas/ResponseFormat" + }, + { + "title": "ResponseFormatList", + "type": "array", + "items": { + "$ref": "#/components/schemas/ResponseFormat" + } + } + ] + }, + "environment": { + "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentConfig" + }, + { + "type": "string" + } + ] + }, + "id": { + "description": "Required. Output only. A unique identifier for the interaction completion.", + "type": "string", + "readOnly": true + }, + "status": { + "description": "Required. Output only. The status of the interaction.", + "type": "string", + "enum": [ + "in_progress", + "requires_action", + "completed", + "failed", + "cancelled", + "incomplete", + "budget_exceeded" + ], + "readOnly": true, + "x-google-enum-deprecated": [ + false, + false, + false, + false, + false, + false, + true + ], + "x-google-enum-descriptions": [ + "The interaction is in progress.", + "The interaction requires action/input from the user.", + "The interaction is completed.", + "The interaction failed.", + "The interaction was cancelled.", + "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + ] + }, + "created": { + "description": "Required. Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "updated": { + "description": "Required. Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "system_instruction": { + "description": "System instruction for the interaction.", + "type": "string" + }, + "tools": { + "description": "A list of tool declarations the model may call during interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Tool" + } + }, + "response_modalities": { + "description": "The requested modalities of the response (TEXT, IMAGE, AUDIO).", + "type": "array", + "deprecated": true, + "items": { + "$ref": "#/components/schemas/ResponseModality" + } + }, + "response_mime_type": { + "description": "The mime type of the response. This is required if response_format is set.", + "type": "string", + "deprecated": true + }, + "previous_interaction_id": { + "description": "The ID of the previous interaction, if any.", + "type": "string" + }, + "environment_id": { + "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", + "type": "string", + "readOnly": true + }, + "service_tier": { + "description": "The service tier for the interaction.", + "$ref": "#/components/schemas/ServiceTier" + }, + "webhook_config": { + "description": "Optional. Webhook configuration for receiving notifications when the\ninteraction completes.", + "$ref": "#/components/schemas/WebhookConfig" + }, + "safety_settings": { + "description": "Safety settings for the interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/SafetySetting" + } + }, + "labels": { + "description": "The labels with user-defined metadata for the request.\n\nLabel keys and values can be no longer than 63 characters\n(Unicode codepoints) and can only contain lowercase letters, numeric\ncharacters, underscores, and dashes. International characters are allowed.\nLabel values are optional. Label keys must start with a letter.", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "stream": { + "description": "Input only. Whether the interaction will be streamed.", + "type": "boolean", + "writeOnly": true + }, + "store": { + "description": "Input only. Whether to store the response and request for later retrieval.", + "type": "boolean", + "writeOnly": true + }, + "background": { + "description": "Input only. Whether to run the model interaction in the background.", + "type": "boolean", + "writeOnly": true + } + }, + "required": [ + "agent", + "input" + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "CreateAgentInteraction" + }, + "CreateEnvironmentRequest": { + "description": "Request for `CreateEnvironment`.", + "type": "object", + "properties": { + "network": { + "description": "Network configuration for the environment.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentNetworkEgressAllowlist" + }, + { + "type": "string", + "enum": [ + "disabled" + ], + "x-google-enum-descriptions": [ + "All network egress is blocked." + ] + } + ] + }, + "sources": { + "description": "Sources to be mounted into the environment.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Source" + } + }, + "from_environment": { + "description": "Optional. The source environment to copy/fork from.\nFormat: `environments/{environment_id}` or `{environment_id}`.\nWhen specified, `sources` and `env` must be empty.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "CreateModelInteractionParams": { + "description": "Parameters for creating model interactions", + "properties": { + "model": { + "description": "The name of the `Model` used for generating the interaction.", + "$ref": "#/components/schemas/ModelOption" + }, + "input": { + "$ref": "#/components/schemas/InteractionsInput" + }, + "generation_config": { + "description": "Input only. Configuration parameters for the model interaction.", + "$ref": "#/components/schemas/GenerationConfig", + "writeOnly": true + }, + "response_format": { + "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", + "oneOf": [ + { + "$ref": "#/components/schemas/ResponseFormat" + }, + { + "title": "ResponseFormatList", + "type": "array", + "items": { + "$ref": "#/components/schemas/ResponseFormat" + } + } + ] + }, + "environment": { + "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentConfig" + }, + { + "type": "string" + } + ] + }, + "id": { + "description": "Required. Output only. A unique identifier for the interaction completion.", + "type": "string", + "readOnly": true + }, + "status": { + "description": "Required. Output only. The status of the interaction.", + "type": "string", + "enum": [ + "in_progress", + "requires_action", + "completed", + "failed", + "cancelled", + "incomplete", + "budget_exceeded" + ], + "readOnly": true, + "x-google-enum-deprecated": [ + false, + false, + false, + false, + false, + false, + true + ], + "x-google-enum-descriptions": [ + "The interaction is in progress.", + "The interaction requires action/input from the user.", + "The interaction is completed.", + "The interaction failed.", + "The interaction was cancelled.", + "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + ] + }, + "created": { + "description": "Required. Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "updated": { + "description": "Required. Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "system_instruction": { + "description": "System instruction for the interaction.", + "type": "string" + }, + "tools": { + "description": "A list of tool declarations the model may call during interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Tool" + } + }, + "response_modalities": { + "description": "The requested modalities of the response (TEXT, IMAGE, AUDIO).", + "type": "array", + "deprecated": true, + "items": { + "$ref": "#/components/schemas/ResponseModality" + } + }, + "response_mime_type": { + "description": "The mime type of the response. This is required if response_format is set.", + "type": "string", + "deprecated": true + }, + "previous_interaction_id": { + "description": "The ID of the previous interaction, if any.", + "type": "string" + }, + "environment_id": { + "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", + "type": "string", + "readOnly": true + }, + "service_tier": { + "description": "The service tier for the interaction.", + "$ref": "#/components/schemas/ServiceTier" + }, + "webhook_config": { + "description": "Optional. Webhook configuration for receiving notifications when the\ninteraction completes.", + "$ref": "#/components/schemas/WebhookConfig" + }, + "safety_settings": { + "description": "Safety settings for the interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/SafetySetting" + } + }, + "labels": { + "description": "The labels with user-defined metadata for the request.\n\nLabel keys and values can be no longer than 63 characters\n(Unicode codepoints) and can only contain lowercase letters, numeric\ncharacters, underscores, and dashes. International characters are allowed.\nLabel values are optional. Label keys must start with a letter.", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "stream": { + "description": "Input only. Whether the interaction will be streamed.", + "type": "boolean", + "writeOnly": true + }, + "store": { + "description": "Input only. Whether to store the response and request for later retrieval.", + "type": "boolean", + "writeOnly": true + }, + "background": { + "description": "Input only. Whether to run the model interaction in the background.", + "type": "boolean", + "writeOnly": true + } + }, + "required": [ + "input", + "model" + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "CreateModelInteraction" + }, + "CreateVoiceRequest": { + "description": "Request message for `VoicesService.CreateVoice`.", + "type": "object", + "properties": { + "voice": { + "description": "Required. The voice to create.\n`voice.model` is optional; if omitted, the service selects the default\nvoice creation model. Output-only fields on `Voice` (`id`, `key`,\n`expire_time`, `usage`) are ignored if set.", + "$ref": "#/components/schemas/Voice" + }, + "store": { + "description": "Optional. Whether the created voice is persisted and managed by Google.\n\n* When `true`, Google stores the voice and returns `Voice.id` (for example,\n `voice_abc123def456`), which can be managed via `GetVoice`, `ListVoices`,\n and `DeleteVoice` and referenced by ID in synthesis requests. Stored\n voices expire after 1 year of inactivity; using a stored voice in speech\n synthesis or as a `base_voice` in `CreateVoice` extends its\n `expire_time`. Projects are subject to a maximum active stored voice\n quota; exceeding the quota returns `RESOURCE_EXHAUSTED`.\n* When `false` (default), the voice is not stored by Google and `Voice.key`\n (for example, `voicekey_...`) is returned for client-side storage and\n synthesis. Optional discovery metadata fields on `voice` are not\n persisted or returned when `store` is `false`.\n* Required to be `true` when `voice.type` is `\"prompted\"`\n (otherwise fails with `INVALID_ARGUMENT`).", + "type": "boolean" + } + }, + "required": [ + "voice" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "CreateWebhookRequest": { + "description": "Request message for WebhookService.CreateWebhook.", + "type": "object", + "properties": { + "webhook": { + "description": "Required. The webhook to create.", + "$ref": "#/components/schemas/Webhook" + } + }, + "required": [ + "webhook" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Credential": { + "description": "Server-managed credential resource stored in Secret Manager.", + "type": "object", + "properties": { + "id": { + "description": "Required. Output only. Identifier. Unique identifier for the credential.", + "type": "string", + "readOnly": true, + "x-google-identifier": true + }, + "type": { + "description": "Required. Output only. The type of credential.", + "type": "string", + "enum": [ + "bearer_token", + "oauth2", + "environment_variable" + ], + "readOnly": true, + "x-google-enum-descriptions": [ + "Static token injected as header. No refresh logic.", + "Auto-refresh expired access tokens using stored refresh token.", + "Environment variable injected into sandbox container." + ] + }, + "status": { + "description": "Output only. Current status of the credential.", + "type": "string", + "enum": [ + "active", + "revoked" + ], + "readOnly": true, + "x-google-enum-descriptions": [ + "The credential is active and valid for use.", + "The credential has been revoked and is no longer valid." + ] + }, + "create_time": { + "description": "Output only. The timestamp when the credential was created.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "update_time": { + "description": "Output only. The timestamp when the credential was last updated.", + "type": "string", + "format": "date-time", + "readOnly": true + } + }, + "required": [ + "id", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "CredentialCreateParams": { + "description": "Represents the fields of a Credential provided on creation.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentVariableConfig" + }, + { + "$ref": "#/components/schemas/HttpBearerConfig" + }, + { + "$ref": "#/components/schemas/OAuth2Config" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "CredentialUpdateParams": { + "description": "Represents the fields of a Credential that can be updated.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentVariableUpdateConfig" + }, + { + "$ref": "#/components/schemas/HttpBearerUpdateConfig" + }, + { + "$ref": "#/components/schemas/OAuth2UpdateConfig" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "DeepResearchAgentConfig": { + "description": "Configuration for the Deep Research agent.", + "type": "object", + "properties": { + "type": { + "const": "deep-research" + }, + "thinking_summaries": { + "description": "Whether to include thought summaries in the response.", + "$ref": "#/components/schemas/ThinkingSummaries" + }, + "visualization": { + "description": "Whether to include visualizations in the response.", + "type": "string", + "enum": [ + "off", + "auto" + ], + "x-google-enum-descriptions": [ + "Do not include visualizations.", + "Automatically include visualizations." + ] + }, + "collaborative_planning": { + "description": "Enables human-in-the-loop planning for the Deep Research agent. If set to\ntrue, the Deep Research agent will provide a research plan in its response.\nThe agent will then proceed only if the user confirms the plan in the next\nturn.", + "type": "boolean" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "DeleteVoiceResponse": { + "description": "Response message for `VoicesService.DeleteVoice`.", + "type": "object", + "x-speakeasy-model-namespace": "interactions" + }, + "DocumentContent": { + "description": "A document content block.", + "type": "object", + "properties": { + "type": { + "const": "document" + }, + "data": { + "description": "The document content.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "description": "The URI of the document.", + "type": "string" + }, + "mime_type": { + "description": "The mime type of the document.", + "type": "string", + "enum": [ + "application/pdf", + "text/csv" + ], + "x-google-enum-descriptions": [ + "PDF document format", + "CSV document format" + ] + } + }, + "required": [ + "type" + ], + "examples": [ + { + "document": { + "summary": "Document", + "value": { + "type": "document", + "data": "BASE64_ENCODED_DOCUMENT", + "mime_type": "application/pdf" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "DocumentDelta": { + "type": "object", + "properties": { + "type": { + "const": "document" + }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, + "mime_type": { + "type": "string", + "enum": [ + "application/pdf", + "text/csv" + ], + "x-google-enum-descriptions": [ + "PDF document format", + "CSV document format" + ] + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "DynamicAgentConfig": { + "description": "Configuration for dynamic agents.", + "type": "object", + "properties": { + "type": { + "const": "dynamic" + } + }, + "required": [ + "type" + ], + "additionalProperties": { + "description": "For agents that are not supported statically in the API definition." + }, + "x-speakeasy-model-namespace": "interactions" + }, + "EgressRule": { + "description": "A single domain allowlist rule with optional header injection.", + "type": "object", + "properties": { + "domain": { + "description": "Domain to allow outbound requests to. Supports wildcards (e.g.\n'*.googleapis.com'). Use '*' to allow all domains.", + "type": "string" + }, + "transform": { + "description": "Headers to inject into requests matching this rule.\nKey: header name (e.g., \"Authorization\").\nValue: header value (e.g., \"Bearer your-token\").", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "credential": { + "description": "Optional. Reference to a server-managed Credential resource by ID.", + "type": "string" + } + }, + "required": [ + "domain" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Empty": { + "description": "A generic empty message that you can re-use to avoid defining duplicated\nempty messages in your APIs. A typical example is to use it as the request\nor the response type of an API method. For instance:\n\n service Foo {\n rpc Bar(google.protobuf.Empty) returns (google.protobuf.Empty);\n }", + "type": "object", + "x-speakeasy-model-namespace": "interactions" + }, + "EnvVar": { + "description": "An environment variable to set in the execution environment.", + "type": "object", + "properties": { + "value": { + "description": "Direct string value for plain environment variables.", + "type": "string" + }, + "credential": { + "description": "Optional reference to a server-managed Credential resource by ID.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "Environment": { + "description": "An execution environment for an agent.", + "type": "object", + "properties": { + "id": { + "description": "Required. Output only. The ID of the environment.", + "type": "string", + "readOnly": true + }, + "network": { + "description": "Network configuration for the environment.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentNetworkEgressAllowlist" + }, + { + "type": "string", + "enum": [ + "disabled" + ], + "x-google-enum-descriptions": [ + "All network egress is blocked." + ] + } + ] + }, + "sources": { + "description": "Sources to be mounted into the environment.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Source" + } + }, + "created": { + "description": "Output only. The time at which the environment was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "updated": { + "description": "Output only. The time at which the environment was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "last_accessed": { + "description": "Output only. The time at which the environment was last accessed in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "status": { + "description": "Output only. The status of the environment container.", + "type": "string", + "enum": [ + "active", + "expired" + ], + "readOnly": true, + "x-google-enum-descriptions": [ + "", + "" + ] + }, + "file_count": { + "description": "Output only. The number of files in the environment, output only.", + "type": "string", + "format": "int64", + "readOnly": true + }, + "size_bytes": { + "description": "Output only. The total size of the environment files in bytes, output only.", + "type": "string", + "format": "int64", + "readOnly": true + } + }, + "required": [ + "id" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "EnvironmentConfig": { + "description": "Configuration for a custom environment.", + "type": "object", + "properties": { + "type": { + "const": "remote" + }, + "network": { + "description": "Network configuration for the environment.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentNetworkEgressAllowlist" + }, + { + "type": "string", + "enum": [ + "disabled" + ], + "x-google-enum-descriptions": [ + "All network egress is blocked." + ] + } + ] + }, + "sources": { + "type": "array", + "items": { + "$ref": "#/components/schemas/Source" + } + }, + "environment_id": { + "description": "Optional. The environment ID for the interaction. If specified, the request will\nupdate the existing environment instead of creating a new one.", + "type": "string" + }, + "env": { + "description": "Environment variables to set in the sandbox environment.", + "oneOf": [ + { + "type": "object", + "additionalProperties": { + "$ref": "#/components/schemas/EnvVar" + } + }, + { + "type": "string" + } + ] + } + }, + "required": [ + "type" + ], + "examples": [ + { + "inline_sources": { + "summary": "Inline Sources", + "value": { + "type": "remote", + "sources": [ + { + "type": "inline", + "content": "You are a data analyst. Always include visualizations and export results as PDF.", + "target": ".agents/AGENTS.md" + }, + { + "type": "inline", + "content": "---\nname: slide-maker\ndescription: Create HTML slide decks\n---\n# Slide Maker\n\nWhen asked to create a presentation:\n1. Analyze the input data\n2. Create an HTML slide deck with reveal.js\n3. Save to /workspace/output/slides.html", + "target": ".agents/skills/slide-maker/SKILL.md" + } + ] + } + } + }, + { + "external_sources": { + "summary": "External Sources", + "value": { + "type": "remote", + "sources": [ + { + "type": "repository", + "source": "https://github.com/my-org/my-skills.git", + "target": ".agents/skills" + }, + { + "type": "gcs", + "source": "gs://my-bucket/my-folder", + "target": "/workspace/data" + } + ] + } + } + }, + { + "network_allowlist": { + "summary": "Network Allowlist", + "value": { + "type": "remote", + "network": { + "allowlist": [ + { + "domain": "pypi.org" + }, + { + "domain": "*.github.com" + } + ] + } + } + } + }, + { + "proxy_credentials": { + "summary": "Proxy Credentials", + "value": { + "type": "remote", + "network": { + "allowlist": [ + { + "domain": "api.github.com", + "transform": { + "Authorization": "Bearer YOUR_GITHUB_TOKEN" + } + } + ] + } + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "EnvironmentFile": { + "description": "Metadata for a file or directory within an environment.", + "type": "object", + "properties": { + "name": { + "description": "Output only. The name of the file or directory (e.g., \"main.py\" or \"src\").", + "type": "string", + "readOnly": true + }, + "path": { + "description": "Output only. The full relative path within the environment\n(e.g., \"workspace/src/main.py\").", + "type": "string", + "readOnly": true + }, + "type": { + "description": "Output only. The type of the entry.", + "type": "string", + "enum": [ + "file", + "directory" + ], + "readOnly": true, + "x-google-enum-descriptions": [ + "A regular file.", + "A directory." + ] + }, + "size_bytes": { + "description": "Output only. The size of the file/directory in bytes.\nNOLINT", + "type": "string", + "format": "int64", + "readOnly": true + }, + "mime_type": { + "description": "Output only. The MIME type of the file (e.g., \"text/python\", \"image/png\").\nEmpty for directories.\nNOLINT", + "type": "string", + "readOnly": true + }, + "created": { + "description": "Output only. The creation time of the file/directory.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "modified": { + "description": "Output only. The modification time of the file/directory.", + "type": "string", + "format": "date-time", + "readOnly": true + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "EnvironmentNetworkEgressAllowlist": { + "description": "Network egress configuration for the environment.", + "type": "object", + "properties": { + "allowlist": { + "description": "List of allowed domains and their configurations. Set to `\"disabled\"`\nto block all network egress.", + "oneOf": [ + { + "type": "array", + "items": { + "$ref": "#/components/schemas/EgressRule" + } + }, + { + "type": "string", + "enum": [ + "disabled" + ], + "x-google-enum-descriptions": [ + "All network egress is blocked." + ] + } + ] + } + }, + "examples": [ + { + "network_allowlist": { + "summary": "Network Allowlist", + "value": { + "allowlist": [ + { + "domain": "github.com", + "transform": [ + { + "Authorization": "Bearer your-token" + } + ] + }, + { + "domain": "*.googleapis.com" + } + ] + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "EnvironmentVariableConfig": { + "description": "Configuration for environment variable credentials.", + "type": "object", + "properties": { + "value": { + "description": "Required. Input only. Secret value of the environment variable. Write-only; never\nreturned in responses.", + "type": "string", + "writeOnly": true + }, + "injection_location": { + "description": "Required. Locations where the environment variable can be injected in\noutgoing HTTP requests. Must contain at least one location.\nAccepts either a single location (e.g. \"header\") or an array of locations.", + "oneOf": [ + { + "$ref": "#/components/schemas/InjectionLocation" + }, + { + "type": "array", + "items": { + "$ref": "#/components/schemas/InjectionLocation" + } + } + ] + }, + "id": { + "type": "string" + }, + "type": { + "const": "environment_variable" + }, + "trusted_domains": { + "description": "Optional. List of domains allowed to receive this environment variable\nvalue in HTTP requests.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "id", + "injection_location", + "type", + "value" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "EnvironmentVariableUpdateConfig": { + "description": "Configuration for updating environment variable credentials.", + "type": "object", + "properties": { + "type": { + "const": "environment_variable" + }, + "value": { + "description": "Optional. Input only. Secret value of the environment variable. Write-only; never\nreturned in responses.", + "type": "string", + "writeOnly": true + }, + "trusted_domains": { + "description": "Optional. List of domains allowed to receive this environment variable\nvalue in HTTP requests.", + "type": "array", + "items": { + "type": "string" + } + }, + "injection_location": { + "description": "Optional. Locations where the environment variable can be injected in\noutgoing HTTP requests.\nAccepts either a single location (e.g. \"header\") or an array of locations.", + "oneOf": [ + { + "$ref": "#/components/schemas/InjectionLocation" + }, + { + "type": "array", + "items": { + "$ref": "#/components/schemas/InjectionLocation" + } + } + ] + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Error": { + "description": "Error message from an interaction.", + "type": "object", + "properties": { + "code": { + "description": "A URI that identifies the error type.", + "type": "string" + }, + "message": { + "description": "A human-readable error message.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ErrorEvent": { + "type": "object", + "properties": { + "event_type": { + "const": "error" + }, + "error": { + "$ref": "#/components/schemas/Error" + }, + "event_id": { + "description": "The event_id token to be used to resume the interaction stream, from\nthis event.", + "type": "string" + } + }, + "required": [ + "event_type" + ], + "examples": [ + { + "error_event": { + "summary": "Error Event", + "value": { + "error": { + "code": "not_found", + "message": "Failed to get completed interaction: Result not found." + }, + "event_type": "error" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "FileCitation": { + "description": "A file citation annotation.", + "type": "object", + "properties": { + "type": { + "const": "file_citation" + }, + "document_uri": { + "description": "The URI of the file.", + "type": "string" + }, + "file_name": { + "description": "The name of the file.", + "type": "string" + }, + "source": { + "description": "Source attributed for a portion of the text.", + "type": "string" + }, + "custom_metadata": { + "description": "User provided metadata about the retrieved context.", + "type": "object", + "additionalProperties": { + "description": "Properties of the object." + } + }, + "page_number": { + "description": "Page number of the cited document, if applicable.", + "type": "integer", + "format": "int32" + }, + "media_id": { + "description": "Media ID in-case of image citations, if applicable.", + "type": "string" + }, + "start_index": { + "description": "Start of segment of the response that is attributed to this source.\n\nIndex indicates the start of the segment, measured in bytes.", + "type": "integer", + "format": "int32" + }, + "end_index": { + "description": "End of the attributed segment, exclusive.", + "type": "integer", + "format": "int32" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "FileSearch": { + "description": "A tool that can be used by the model to search files.", + "type": "object", + "properties": { + "type": { + "const": "file_search" + }, + "file_search_store_names": { + "description": "The file search store names to search.", + "type": "array", + "items": { + "type": "string" + } + }, + "top_k": { + "description": "The number of semantic retrieval chunks to retrieve.", + "type": "integer", + "format": "int32" + }, + "metadata_filter": { + "description": "Metadata filter to apply to the semantic retrieval documents and chunks.", + "type": "string" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-codeSamples": [ + { + "label": "file_search", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"tools\": [{\n \"type\": \"file_search\",\n \"file_search_store_names\": [\"fileSearchStores/m64d1sevsr4y-xfyawui3fxqg\"]\n }],\n \"input\": \"Who is the author of the book?\"\n }'\n" + }, + { + "label": "file_search", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\n\n# Create a file search store so we have a valid one to use.\nstore = client.file_search_stores.create()\n\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n tools=[\n {\"type\": \"file_search\", \"file_search_store_names\": [store.name]}\n ],\n input=\"What documents are available?\",\n)\nprint(response.output_text)\n" + }, + { + "label": "file_search", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\n\n// Create a file search store so we have a valid one to use.\nconst store = await ai.fileSearchStores.create({});\nif (!store.name) {\n throw new Error('Store creation failed: Name is undefined');\n}\n\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n tools: [\n {\n type: 'file_search',\n file_search_store_names: [store.name],\n },\n ],\n input: 'What documents are available?',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "file_search", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.types.CreateFileSearchStoreConfig;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.FileSearch;\nimport com.google.genai.types.FileSearchStore;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\n\n// Create a file search store so we have a valid one to use.\nFileSearchStore store =\n client.fileSearchStores.create(CreateFileSearchStoreConfig.builder().build());\nString storeName = store.name().orElseThrow();\n\nFileSearch tool = FileSearch.builder().fileSearchStoreNames(List.of(storeName)).build();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .tools(List.of(tool))\n .input(InteractionsInput.of(\"What documents are available?\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "file_search", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n tool := gaos_interactions.NewTool(gaos_interactions.FileSearch{\n FileSearchStoreNames: []string{store.Name},\n })\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"What documents are available?\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + } + ] + }, + "FileSearchCallDelta": { + "type": "object", + "properties": { + "type": { + "const": "file_search_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "FileSearchCallStep": { + "description": "File Search call step.", + "type": "object", + "properties": { + "id": { + "description": "Required. A unique ID for this specific tool call.", + "type": "string" + }, + "type": { + "const": "file_search_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "id", + "type" + ], + "examples": [ + { + "file_search_call": { + "summary": "FileSearchCallStep", + "value": { + "type": "file_search_call", + "id": "file_call_88192" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "FileSearchResult": { + "description": "The result of the File Search.", + "type": "object", + "x-speakeasy-model-namespace": "interactions" + }, + "FileSearchResultDelta": { + "type": "object", + "properties": { + "result": { + "type": "array", + "items": { + "$ref": "#/components/schemas/FileSearchResult" + } + }, + "type": { + "const": "file_search_result" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "result", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "FileSearchResultStep": { + "description": "File Search result step.", + "type": "object", + "properties": { + "call_id": { + "description": "Required. ID to match the ID from the function call block.", + "type": "string" + }, + "type": { + "const": "file_search_result" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "call_id", + "type" + ], + "examples": [ + { + "file_search_result": { + "summary": "FileSearchResultStep", + "value": { + "type": "file_search_result", + "call_id": "file_call_88192" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Function": { + "description": "A tool that can be used by the model.", + "type": "object", + "properties": { + "type": { + "const": "function" + }, + "name": { + "description": "The name of the function.", + "type": "string" + }, + "description": { + "description": "A description of the function.", + "type": "string" + }, + "parameters": { + "description": "The JSON Schema for the function's parameters.\n\nSupports a subset of JSON Schema. `minLength`, `maxLength`,\n`minProperties`, `maxProperties`, `uniqueItems`, `contains`,\n`minContains`, and `maxContains` are ignored and not enforced." + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-codeSamples": [ + { + "label": "function_calling", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"tools\": [{\n \"type\": \"function\",\n \"name\": \"get_weather\",\n \"description\": \"Get the current weather in a given location\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"location\": {\n \"type\": \"string\",\n \"description\": \"The city and state, e.g. San Francisco, CA\"\n }\n },\n \"required\": [\"location\"]\n }\n }],\n \"input\": \"What is the weather like in Boston, MA?\"\n }'\n" + }, + { + "label": "function_calling", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n tools=[{\n \"type\": \"function\",\n \"name\": \"get_weather\",\n \"description\": \"Get the current weather in a given location\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"location\": {\n \"type\": \"string\",\n \"description\": (\n \"The city and state, e.g. San Francisco, CA\"\n ),\n }\n },\n \"required\": [\"location\"],\n },\n }],\n input=\"What is the weather like in Boston?\",\n)\nprint(response.steps[-1])\n" + }, + { + "label": "function_calling", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n tools: [\n {\n type: 'function',\n name: 'get_weather',\n description: 'Get the current weather in a given location',\n parameters: {\n type: 'object',\n properties: {\n location: {\n type: 'string',\n description: 'The city and state, e.g. San Francisco, CA',\n },\n },\n required: ['location'],\n },\n },\n ],\n input: 'What is the weather like in Boston?',\n});\nconsole.log(interaction.steps.at(-1));\n" + }, + { + "label": "function_calling", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Function;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.Step;\n\nClient client = new Client();\nMap parameters =\n Map.of(\n \"type\", \"object\",\n \"properties\",\n Map.of(\n \"location\",\n Map.of(\n \"type\", \"string\",\n \"description\", \"The city and state, e.g. San Francisco, CA\")),\n \"required\", List.of(\"location\"));\nFunction functionTool =\n Function.builder()\n .name(\"get_weather\")\n .description(\"Get the current weather in a given location\")\n .parameters(parameters)\n .build();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .tools(List.of(functionTool))\n .input(InteractionsInput.of(\"What is the weather like in Boston?\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nList steps = interaction.steps().orElse(List.of());\nif (!steps.isEmpty()) {\n System.out.println(steps.get(steps.size() - 1));\n}\n" + }, + { + "label": "function_calling", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n function := gaos_interactions.Function{\n Name: genai.Ptr(\"get_weather\"),\n Description: genai.Ptr(\"Get the current weather in a given location\"),\n Parameters: map[string]any{\n \"type\": \"object\",\n \"properties\": map[string]any{\n \"location\": map[string]any{\n \"type\": \"string\",\n \"description\": \"The city and state, e.g. San Francisco, CA\",\n },\n },\n \"required\": []string{\"location\"},\n },\n }\n\n tool := gaos_interactions.NewTool(function)\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"What is the weather like in Boston, MA?\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil {\n for _, step := range res.Interaction.Steps {\n if step.FunctionCallStep != nil {\n fmt.Println(\"Function Call:\", step.FunctionCallStep.Name)\n fmt.Println(\"Arguments:\", step.FunctionCallStep.Arguments)\n }\n }\n }\n}\n" + } + ] + }, + "FunctionCallStep": { + "description": "A function tool call step.", + "type": "object", + "properties": { + "name": { + "description": "Required. The name of the tool to call.", + "type": "string" + }, + "arguments": { + "description": "Required. The arguments to pass to the function.", + "type": "object", + "additionalProperties": { + "description": "Properties of the object." + } + }, + "id": { + "description": "Required. A unique ID for this specific tool call.", + "type": "string" + }, + "type": { + "const": "function_call" + } + }, + "required": [ + "arguments", + "id", + "name", + "type" + ], + "examples": [ + { + "function_call": { + "summary": "FunctionCallStep", + "value": { + "name": "get_weather", + "type": "function_call", + "arguments": { + "location": "Boston, MA" + }, + "id": "call_98231" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "FunctionResultDelta": { + "type": "object", + "properties": { + "result": { + "oneOf": [ + { + "type": "array", + "items": { + "oneOf": [ + { + "$ref": "#/components/schemas/ImageContent" + }, + { + "$ref": "#/components/schemas/TextContent" + } + ] + } + }, + { + "type": "object" + }, + { + "type": "string" + } + ] + }, + "type": { + "const": "function_result" + }, + "name": { + "type": "string" + }, + "is_error": { + "type": "boolean" + } + }, + "required": [ + "result", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "FunctionResultStep": { + "description": "Result of a function tool call.", + "type": "object", + "properties": { + "result": { + "description": "Required. The result of the tool call.", + "oneOf": [ + { + "type": "array", + "items": { + "oneOf": [ + { + "$ref": "#/components/schemas/ImageContent" + }, + { + "$ref": "#/components/schemas/TextContent" + } + ] + } + }, + { + "type": "object" + }, + { + "type": "string" + } + ] + }, + "call_id": { + "description": "Required. ID to match the ID from the function call block.", + "type": "string" + }, + "type": { + "const": "function_result" + }, + "name": { + "description": "The name of the tool that was called.", + "type": "string" + }, + "is_error": { + "description": "Whether the tool call resulted in an error.", + "type": "boolean" + } + }, + "required": [ + "call_id", + "result", + "type" + ], + "examples": [ + { + "function_result": { + "summary": "FunctionResultStep", + "value": { + "name": "get_weather", + "type": "function_result", + "call_id": "call_98231", + "result": [ + { + "type": "text", + "text": "{\"weather\":\"sunny\"}" + } + ] + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GenerationConfig": { + "description": "Configuration parameters for model interactions.", + "type": "object", + "properties": { + "tool_choice": { + "description": "The tool choice configuration.", + "oneOf": [ + { + "$ref": "#/components/schemas/ToolChoiceConfig" + }, + { + "type": "string", + "enum": [ + "auto", + "any", + "none", + "validated" + ], + "x-google-enum-descriptions": [ + "Auto tool choice.", + "Any tool choice.", + "No tool choice.", + "Validated tool choice." + ], + "x-speakeasy-model-namespace": "interactions" + } + ] + }, + "temperature": { + "description": "Controls the randomness of the output.", + "type": "number", + "deprecated": true, + "format": "float" + }, + "top_p": { + "description": "The maximum cumulative probability of tokens to consider when sampling.", + "type": "number", + "deprecated": true, + "format": "float" + }, + "seed": { + "description": "Seed used in decoding for reproducibility.", + "type": "integer", + "format": "int32" + }, + "stop_sequences": { + "description": "A list of character sequences that will stop output interaction.", + "type": "array", + "items": { + "type": "string" + } + }, + "thinking_level": { + "description": "The level of thought tokens that the model should generate.", + "$ref": "#/components/schemas/ThinkingLevel" + }, + "thinking_summaries": { + "description": "Whether to include thought summaries in the response.", + "$ref": "#/components/schemas/ThinkingSummaries" + }, + "max_output_tokens": { + "description": "The maximum number of tokens to include in the response.", + "type": "integer", + "format": "int32" + }, + "image_config": { + "description": "Configuration for image interaction.", + "$ref": "#/components/schemas/ImageConfig", + "deprecated": true + }, + "video_config": { + "description": "Configuration for video generation.", + "$ref": "#/components/schemas/VideoConfig" + }, + "transcription_config": { + "description": "Optional. Configuration for speech recognition (transcription). If present, ASR is\nenabled.", + "$ref": "#/components/schemas/TranscriptionConfig" + }, + "structured_speech_config": { + "description": "Speech and multi-speaker configuration.", + "oneOf": [ + { + "$ref": "#/components/schemas/SpeakerConfig" + }, + { + "type": "array", + "items": { + "$ref": "#/components/schemas/SpeechConfig" + } + } + ] + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "GetEnvironmentFilesResponse": { + "description": "Response for `GetEnvironmentFiles`.", + "type": "object", + "properties": { + "files": { + "description": "If the requested path is a directory, this contains its contents.\nIf the requested path is a file, this contains a single entry with the\nfile's metadata.\nIf alt=media was specified, this is empty (content is served via `blob`).", + "type": "array", + "items": { + "$ref": "#/components/schemas/EnvironmentFile" + } + }, + "next_page_token": { + "description": "Pagination token for directory listing.\nNOLINT", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleMaps": { + "description": "A tool that can be used by the model to call Google Maps.", + "type": "object", + "properties": { + "type": { + "const": "google_maps" + }, + "enable_widget": { + "description": "Whether to return a widget context token in the tool call result of the\nresponse.", + "type": "boolean" + }, + "latitude": { + "description": "The latitude of the user's location.", + "type": "number", + "format": "double" + }, + "longitude": { + "description": "The longitude of the user's location.", + "type": "number", + "format": "double" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-codeSamples": [ + { + "label": "google_maps", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"tools\": [{\n \"type\": \"google_maps\",\n \"latitude\": 37.7749,\n \"longitude\": -122.4194\n }],\n \"input\": \"What is the best food near me?\"\n }'\n" + }, + { + "label": "google_maps", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n tools=[\n {\"type\": \"google_maps\", \"latitude\": 37.7749, \"longitude\": -122.4194}\n ],\n input=\"What is the best food near me?\",\n)\nprint(response.output_text)\n" + }, + { + "label": "google_maps", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n tools: [\n {\n type: 'google_maps',\n latitude: 37.7749,\n longitude: -122.4194,\n },\n ],\n input: 'What is the best food near me?',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "google_maps", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.GoogleMaps;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nGoogleMaps tool = GoogleMaps.builder().latitude(37.7749).longitude(-122.4194).build();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .tools(List.of(tool))\n .input(InteractionsInput.of(\"What is the best food near me?\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "google_maps", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n tool := gaos_interactions.NewTool(gaos_interactions.GoogleMaps{\n Latitude: genai.Ptr(37.7749),\n Longitude: genai.Ptr(-122.4194),\n })\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"What is the best food near me?\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + } + ] + }, + "GoogleMapsCallArguments": { + "description": "The arguments to pass to the Google Maps tool.", + "type": "object", + "properties": { + "queries": { + "description": "The queries to be executed.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleMapsCallDelta": { + "type": "object", + "properties": { + "type": { + "const": "google_maps_call" + }, + "arguments": { + "description": "The arguments to pass to the Google Maps tool.", + "$ref": "#/components/schemas/GoogleMapsCallArguments" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleMapsCallStep": { + "description": "Google Maps call step.", + "type": "object", + "properties": { + "id": { + "description": "Required. A unique ID for this specific tool call.", + "type": "string" + }, + "type": { + "const": "google_maps_call" + }, + "arguments": { + "description": "The arguments to pass to the Google Maps tool.", + "$ref": "#/components/schemas/GoogleMapsCallStepArguments" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "id", + "type" + ], + "examples": [ + { + "google_maps_call": { + "summary": "GoogleMapsCallStep", + "value": { + "type": "google_maps_call", + "arguments": { + "latitude": 37.7749, + "longitude": -122.4194 + }, + "id": "maps_call_39201" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleMapsCallStepArguments": { + "description": "The arguments to pass to the Google Maps tool.", + "type": "object", + "properties": { + "queries": { + "description": "The queries to be executed.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "GoogleMapsCallArguments" + }, + "GoogleMapsResult": { + "description": "The result of the Google Maps.", + "type": "object", + "properties": { + "places": { + "description": "The places that were found.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Places" + } + }, + "widget_context_token": { + "description": "Resource name of the Google Maps widget context token.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleMapsResultDelta": { + "type": "object", + "properties": { + "type": { + "const": "google_maps_result" + }, + "result": { + "description": "The results of the Google Maps.", + "type": "array", + "items": { + "$ref": "#/components/schemas/GoogleMapsResult" + } + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleMapsResultItem": { + "description": "The result of the Google Maps.", + "type": "object", + "properties": { + "places": { + "type": "array", + "items": { + "$ref": "#/components/schemas/GoogleMapsResultPlaces" + } + }, + "widget_context_token": { + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "GoogleMapsResult" + }, + "GoogleMapsResultPlaces": { + "type": "object", + "properties": { + "place_id": { + "type": "string" + }, + "name": { + "type": "string" + }, + "url": { + "type": "string" + }, + "review_snippets": { + "type": "array", + "items": { + "$ref": "#/components/schemas/ReviewSnippet" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleMapsResultStep": { + "description": "Google Maps result step.", + "type": "object", + "properties": { + "result": { + "type": "array", + "items": { + "$ref": "#/components/schemas/GoogleMapsResultItem" + } + }, + "call_id": { + "description": "Required. ID to match the ID from the function call block.", + "type": "string" + }, + "type": { + "const": "google_maps_result" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "call_id", + "result", + "type" + ], + "examples": [ + { + "google_maps_result": { + "summary": "GoogleMapsResultStep", + "value": { + "type": "google_maps_result", + "call_id": "maps_call_39201", + "result": [ + { + "name": "Golden Gate Park", + "place_id": "ChIJIQBpAG2ahYAR9R7bNdTLg8M", + "rating": 4.8 + } + ] + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleSearch": { + "description": "A tool that can be used by the model to search Google.", + "type": "object", + "properties": { + "type": { + "const": "google_search" + }, + "search_types": { + "description": "The types of search grounding to enable.", + "type": "array", + "items": { + "type": "string", + "enum": [ + "web_search", + "image_search" + ], + "x-google-enum-descriptions": [ + "Setting this field enables web search. Only text results are returned.", + "Setting this field enables image search. Image bytes are returned." + ] + } + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-codeSamples": [ + { + "label": "google_search", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"tools\": [{\n \"type\": \"google_search\"\n }],\n \"input\": \"Who is the current president of France?\"\n }'\n" + }, + { + "label": "google_search", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n tools=[{\"type\": \"google_search\"}],\n input=\"Who is the current president of France?\",\n)\nprint(response.output_text)\n" + }, + { + "label": "google_search", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n tools: [{type: 'google_search'}],\n input: 'Who is the current president of France?',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "google_search", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.GoogleSearch;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .tools(List.of(new GoogleSearch()))\n .input(InteractionsInput.of(\"Who is the current president of France?\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "google_search", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n tool := gaos_interactions.NewTool(gaos_interactions.GoogleSearch{})\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Who is the current president of France?\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + } + ] + }, + "GoogleSearchCallArguments": { + "description": "The arguments to pass to Google Search.", + "type": "object", + "properties": { + "queries": { + "description": "Web search queries for the following-up web search.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleSearchCallDelta": { + "type": "object", + "properties": { + "arguments": { + "$ref": "#/components/schemas/GoogleSearchCallArguments" + }, + "type": { + "const": "google_search_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "arguments", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleSearchCallStep": { + "description": "Google Search call step.", + "type": "object", + "properties": { + "arguments": { + "description": "Required. The arguments to pass to Google Search.", + "$ref": "#/components/schemas/GoogleSearchCallStepArguments" + }, + "id": { + "description": "Required. A unique ID for this specific tool call.", + "type": "string" + }, + "type": { + "const": "google_search_call" + }, + "search_type": { + "description": "The type of search grounding enabled.", + "type": "string", + "enum": [ + "web_search", + "image_search" + ], + "x-google-enum-descriptions": [ + "Setting this field enables web search. Only text results are returned.", + "Setting this field enables image search. Image bytes are returned." + ] + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "arguments", + "id", + "type" + ], + "examples": [ + { + "google_search_call": { + "summary": "GoogleSearchCallStep", + "value": { + "type": "google_search_call", + "arguments": { + "query": "Who won the men's 100m in Paris 2024?" + }, + "id": "search_call_19201" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleSearchCallStepArguments": { + "description": "The arguments to pass to Google Search.", + "type": "object", + "properties": { + "queries": { + "description": "Web search queries for the following-up web search.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "GoogleSearchCallArguments" + }, + "GoogleSearchResult": { + "description": "The result of the Google Search.", + "type": "object", + "properties": { + "search_suggestions": { + "description": "Web content snippet that can be embedded in a web page or an app webview.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleSearchResultDelta": { + "type": "object", + "properties": { + "result": { + "type": "array", + "items": { + "$ref": "#/components/schemas/GoogleSearchResult" + } + }, + "type": { + "const": "google_search_result" + }, + "is_error": { + "type": "boolean" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "result", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GoogleSearchResultItem": { + "description": "The result of the Google Search.", + "type": "object", + "properties": { + "search_suggestions": { + "description": "Web content snippet that can be embedded in a web page or an app webview.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "GoogleSearchResult" + }, + "GoogleSearchResultStep": { + "description": "Google Search result step.", + "type": "object", + "properties": { + "result": { + "description": "Required. The results of the Google Search.", + "type": "array", + "items": { + "$ref": "#/components/schemas/GoogleSearchResultItem" + } + }, + "call_id": { + "description": "Required. ID to match the ID from the function call block.", + "type": "string" + }, + "type": { + "const": "google_search_result" + }, + "is_error": { + "description": "Whether the Google Search resulted in an error.", + "type": "boolean" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "call_id", + "result", + "type" + ], + "examples": [ + { + "google_search_result": { + "summary": "GoogleSearchResultStep", + "value": { + "type": "google_search_result", + "call_id": "search_call_19201", + "result": [ + { + "title": "Paris 2024 Olympics: Noah Lyles wins men's 100m gold", + "url": "https://olympics.com/en/news/paris-2024-noah-lyles-wins-mens-100m-gold", + "snippet": "American Noah Lyles won the Olympic men's 100m gold medal in a photo finish." + } + ] + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "GroundingToolCount": { + "description": "The number of grounding tool counts.", + "type": "object", + "properties": { + "type": { + "description": "The grounding tool type associated with the count.", + "type": "string", + "enum": [ + "google_search", + "google_maps" + ], + "x-google-enum-descriptions": [ + "Grounding with Google Web Search and Image Search, & Web Grounding\nfor Enterprise.", + "Grounding with Google Maps." + ] + }, + "count": { + "description": "The number of grounding tool counts.", + "type": "integer", + "format": "int32" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "HarmCategory": { + "type": "string", + "enum": [ + "hate_speech", + "dangerous_content", + "harassment", + "sexually_explicit", + "jailbreak" + ], + "x-google-enum-descriptions": [ + "Content that promotes violence or incites hatred against individuals or\ngroups based on certain attributes.", + "Content that promotes, facilitates, or enables dangerous activities.", + "Abusive, threatening, or content intended to bully, torment, or ridicule.", + "Content that contains sexually explicit material.", + "Prompts designed to bypass safety filters." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "HttpBearerConfig": { + "description": "Configuration for HTTP Bearer token credentials.", + "type": "object", + "properties": { + "token": { + "description": "Required. Input only. The static bearer token. Write-only; never returned in responses.", + "type": "string", + "writeOnly": true + }, + "id": { + "type": "string" + }, + "type": { + "const": "bearer_token" + }, + "header_name": { + "description": "Optional. Header name to inject the token into. Defaults to\n'Authorization'.", + "type": "string" + }, + "prefix": { + "description": "Optional. Prefix to prepend to the token. Defaults to 'Bearer'. Set to ''\nfor no prefix.", + "type": "string" + } + }, + "required": [ + "id", + "token", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "HttpBearerUpdateConfig": { + "description": "Configuration for updating HTTP Bearer token credentials.", + "type": "object", + "properties": { + "type": { + "const": "bearer_token" + }, + "token": { + "description": "Optional. Input only. The static bearer token. Write-only; never returned in responses.", + "type": "string", + "writeOnly": true + }, + "header_name": { + "description": "Optional. Header name to inject the token into. Defaults to\n'Authorization'.", + "type": "string" + }, + "prefix": { + "description": "Optional. Prefix to prepend to the token. Defaults to 'Bearer'. Set to ''\nfor no prefix.", + "type": "string" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "HttpBody": { + "description": "Message that represents an arbitrary HTTP body. It should only be used for\npayload formats that can't be represented as JSON, such as raw binary or\nan HTML page.\n\n\nThis message can be used both in streaming and non-streaming API methods in\nthe request as well as the response.\n\nIt can be used as a top-level request field, which is convenient if one\nwants to extract parameters from either the URL or HTTP template into the\nrequest fields and also want access to the raw HTTP body.\n\nExample:\n\n message GetResourceRequest {\n // A unique request id.\n string request_id = 1;\n\n // The raw HTTP body is bound to this field.\n google.api.HttpBody http_body = 2;\n\n }\n\n service ResourceService {\n rpc GetResource(GetResourceRequest)\n returns (google.api.HttpBody);\n rpc UpdateResource(google.api.HttpBody)\n returns (google.protobuf.Empty);\n\n }\n\nExample with streaming methods:\n\n service CaldavService {\n rpc GetCalendar(stream google.api.HttpBody)\n returns (stream google.api.HttpBody);\n rpc UpdateCalendar(stream google.api.HttpBody)\n returns (stream google.api.HttpBody);\n\n }\n\nUse of this type only changes how the request and response bodies are\nhandled, all other features will continue to work unchanged.", + "type": "object", + "properties": { + "content_type": { + "description": "The HTTP Content-Type header value specifying the content type of the body.", + "type": "string" + }, + "data": { + "description": "The HTTP request/response body as raw binary.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "extensions": { + "description": "Application specific response metadata. Must be set in the first response\nfor streaming APIs.", + "type": "array", + "items": { + "type": "object", + "additionalProperties": { + "description": "Properties of the object. Contains field @type with type URL." + } + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ImageConfig": { + "description": "The configuration for image interaction.", + "type": "object", + "properties": { + "aspect_ratio": { + "description": "The aspect ratio of the image to generate. Supported aspect ratios: 1:1,\n2:3, 3:2, 3:4, 4:3, 9:16, 16:9, 21:9.\n\nIf not specified, the model will choose a default aspect ratio based on any\nreference images provided.", + "type": "string" + }, + "image_size": { + "description": "Specifies the size of generated images. Supported values are `1K`, `2K`,\n`4K`. If not specified, the model will use default value `1K`.", + "type": "string" + } + }, + "deprecated": true, + "x-speakeasy-model-namespace": "interactions" + }, + "ImageContent": { + "description": "An image content block.", + "type": "object", + "properties": { + "type": { + "const": "image" + }, + "data": { + "description": "The image content.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "description": "The URI of the image.", + "type": "string" + }, + "resolution": { + "description": "The resolution of the media.", + "$ref": "#/components/schemas/MediaResolution" + }, + "mime_type": { + "description": "The mime type of the image.", + "type": "string", + "enum": [ + "image/png", + "image/jpeg", + "image/webp", + "image/heic", + "image/heif", + "image/gif", + "image/bmp", + "image/tiff" + ], + "x-google-enum-descriptions": [ + "PNG image format", + "JPEG image format", + "WebP image format", + "HEIC image format", + "HEIF image format", + "GIF image format", + "BMP image format", + "TIFF image format" + ] + } + }, + "required": [ + "type" + ], + "examples": [ + { + "image": { + "summary": "Image", + "value": { + "type": "image", + "data": "BASE64_ENCODED_IMAGE", + "mime_type": "image/png" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ImageDelta": { + "type": "object", + "properties": { + "type": { + "const": "image" + }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, + "mime_type": { + "type": "string", + "enum": [ + "image/png", + "image/jpeg", + "image/webp", + "image/heic", + "image/heif", + "image/gif", + "image/bmp", + "image/tiff" + ], + "x-google-enum-descriptions": [ + "PNG image format", + "JPEG image format", + "WebP image format", + "HEIC image format", + "HEIF image format", + "GIF image format", + "BMP image format", + "TIFF image format" + ] + }, + "resolution": { + "description": "The resolution of the media.", + "$ref": "#/components/schemas/MediaResolution" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ImageResponseFormat": { + "description": "Configuration for image output format.", + "type": "object", + "properties": { + "type": { + "const": "image" + }, + "mime_type": { + "description": "The MIME type of the image output.", + "type": "string", + "enum": [ + "image/jpeg" + ], + "x-google-enum-descriptions": [ + "JPEG image format." + ] + }, + "delivery": { + "description": "The delivery mode for the image output.", + "type": "string", + "enum": [ + "inline", + "uri" + ], + "x-google-enum-descriptions": [ + "Image data is returned inline in the response.", + "Image data is returned as a URI." + ] + }, + "aspect_ratio": { + "description": "The aspect ratio for the image output.", + "type": "string", + "enum": [ + "1:1", + "2:3", + "3:2", + "3:4", + "4:3", + "4:5", + "5:4", + "9:16", + "16:9", + "21:9", + "1:8", + "8:1", + "1:4", + "4:1" + ], + "x-google-enum-descriptions": [ + "1:1 aspect ratio.", + "2:3 aspect ratio.", + "3:2 aspect ratio.", + "3:4 aspect ratio.", + "4:3 aspect ratio.", + "4:5 aspect ratio.", + "5:4 aspect ratio.", + "9:16 aspect ratio.", + "16:9 aspect ratio.", + "21:9 aspect ratio.", + "1:8 aspect ratio.", + "8:1 aspect ratio.", + "1:4 aspect ratio.", + "4:1 aspect ratio." + ] + }, + "image_size": { + "description": "The size of the image output.", + "type": "string", + "enum": [ + "512", + "1K", + "2K", + "4K" + ], + "x-google-enum-descriptions": [ + "512px image size.", + "1K image size.", + "2K image size.", + "4K image size." + ] + } + }, + "required": [ + "type" + ], + "examples": [ + { + "image_response_format": { + "summary": "Image Output", + "value": { + "type": "image", + "aspect_ratio": "16:9", + "image_size": "1K", + "mime_type": "image/jpeg" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "InjectionLocation": { + "type": "string", + "enum": [ + "header", + "query", + "body" + ], + "x-google-enum-descriptions": [ + "Injected into HTTP request headers.", + "Injected into HTTP URL query parameters.", + "Injected into HTTP request body." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Interaction": { + "description": "The Interaction resource.", + "properties": { + "id": { + "description": "Required. Output only. A unique identifier for the interaction completion.", + "type": "string", + "default": "", + "readOnly": true + }, + "status": { + "description": "Required. Output only. The status of the interaction.", + "type": "string", + "enum": [ + "in_progress", + "requires_action", + "completed", + "failed", + "cancelled", + "incomplete", + "budget_exceeded" + ], + "readOnly": true, + "x-google-enum-deprecated": [ + false, + false, + false, + false, + false, + false, + true + ], + "x-google-enum-descriptions": [ + "The interaction is in progress.", + "The interaction requires action/input from the user.", + "The interaction is completed.", + "The interaction failed.", + "The interaction was cancelled.", + "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + ] + }, + "created": { + "description": "Required. Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "updated": { + "description": "Required. Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", + "type": "string", + "readOnly": true + }, + "steps": { + "description": "Required. Output only. The steps that make up the interaction, when included in the response.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Step" + }, + "readOnly": true + }, + "model": { + "description": "The name of the `Model` used for generating the interaction.", + "$ref": "#/components/schemas/ModelOption" + }, + "generation_config": { + "description": "Input only. Configuration parameters for the model interaction.", + "$ref": "#/components/schemas/GenerationConfig", + "writeOnly": true + }, + "input": { + "$ref": "#/components/schemas/InteractionsInput" + }, + "response_format": { + "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", + "oneOf": [ + { + "$ref": "#/components/schemas/ResponseFormat" + }, + { + "title": "ResponseFormatList", + "type": "array", + "items": { + "$ref": "#/components/schemas/ResponseFormat" + } + } + ] + }, + "environment": { + "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentConfig" + }, + { + "type": "string" + } + ] + }, + "system_instruction": { + "description": "System instruction for the interaction.", + "type": "string" + }, + "tools": { + "description": "A list of tool declarations the model may call during interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Tool" + } + }, + "usage": { + "description": "Output only. Statistics on the interaction request's token usage.", + "$ref": "#/components/schemas/Usage", + "readOnly": true + }, + "response_modalities": { + "description": "The requested modalities of the response (TEXT, IMAGE, AUDIO).", + "type": "array", + "deprecated": true, + "items": { + "$ref": "#/components/schemas/ResponseModality" + } + }, + "response_mime_type": { + "description": "The mime type of the response. This is required if response_format is set.", + "type": "string", + "deprecated": true + }, + "previous_interaction_id": { + "description": "The ID of the previous interaction, if any.", + "type": "string" + }, + "environment_id": { + "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", + "type": "string", + "readOnly": true + }, + "service_tier": { + "description": "The service tier for the interaction.", + "$ref": "#/components/schemas/ServiceTier" + }, + "webhook_config": { + "description": "Optional. Webhook configuration for receiving notifications when the\ninteraction completes.", + "$ref": "#/components/schemas/WebhookConfig" + }, + "safety_settings": { + "description": "Safety settings for the interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/SafetySetting" + } + }, + "labels": { + "description": "The labels with user-defined metadata for the request.\n\nLabel keys and values can be no longer than 63 characters\n(Unicode codepoints) and can only contain lowercase letters, numeric\ncharacters, underscores, and dashes. International characters are allowed.\nLabel values are optional. Label keys must start with a letter.", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "errors": { + "description": "Output only. Diagnostic faults / platform errors recorded on the interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Error" + }, + "readOnly": true + }, + "stream": { + "description": "Input only. Whether the interaction will be streamed.", + "type": "boolean", + "writeOnly": true + }, + "store": { + "description": "Input only. Whether to store the response and request for later retrieval.", + "type": "boolean", + "writeOnly": true + }, + "background": { + "description": "Input only. Whether to run the model interaction in the background.", + "type": "boolean", + "writeOnly": true + }, + "agent": { + "description": "The name of the `Agent` used for generating the interaction.", + "$ref": "#/components/schemas/AgentOption" + }, + "agent_config": { + "description": "Configuration parameters for the agent interaction.", + "oneOf": [ + { + "$ref": "#/components/schemas/AntigravityAgentConfig" + }, + { + "$ref": "#/components/schemas/DeepResearchAgentConfig" + }, + { + "$ref": "#/components/schemas/DynamicAgentConfig" + } + ] + }, + "max_total_tokens": { + "description": "Max total tokens for the agent run.", + "type": "string", + "format": "int64" + } + }, + "required": [ + "status", + "created", + "id", + "steps", + "updated" + ], + "example": { + "created": "2025-12-04T15:01:45Z", + "id": "v1_ChdXS0l4YWZXTk9xbk0xZThQczhEcmlROBIXV0tJeGFmV05PcW5NMWU4UHM4RHJpUTg", + "model": "gemini-3.6-flash", + "object": "interaction", + "status": "completed", + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "Hello! I'm doing well, functioning as expected. Thank you for asking! How are you doing today?" + } + ] + } + ], + "updated": "2025-12-04T15:01:45Z", + "usage": { + "input_tokens_by_modality": [ + { + "modality": "text", + "tokens": 7 + } + ], + "total_cached_tokens": 0, + "total_input_tokens": 7, + "total_output_tokens": 23, + "total_thought_tokens": 49, + "total_tokens": 79, + "total_tool_use_tokens": 0 + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "InteractionCompletedEvent": { + "description": "Signals that the Interaction completed. Sent when the Interaction receives\nComplete/Cancel or naturally terminates. No more input can be sent to the\nInteraction after this.", + "type": "object", + "properties": { + "interaction": { + "description": "Required. Partial completed interaction resource emitted at the end of the stream.", + "$ref": "#/components/schemas/InteractionSseEventInteraction" + }, + "event_type": { + "const": "interaction.completed" + }, + "event_id": { + "description": "The event_id token to be used to resume the interaction stream, from\nthis event.", + "type": "string" + } + }, + "required": [ + "event_type", + "interaction" + ], + "examples": [ + { + "interaction_completed": { + "summary": "Interaction Completed", + "value": { + "event_id": "evt_123", + "event_type": "interaction.completed", + "interaction": { + "created": "2025-12-04T15:01:45Z", + "id": "v1_ChdXS0l4YWZXTk9xbk0xZThQczhEcmlROBIXV0tJeGFmV05PcW5NMWU4UHM4RHJpUTg", + "model": "gemini-3.6-flash", + "status": "completed", + "updated": "2025-12-04T15:01:45Z" + } + } + } + }, + { + "interaction_completed": { + "summary": "Interaction Completed", + "value": { + "event_id": "evt_123", + "event_type": "interaction.completed", + "interaction": { + "created": "2025-12-04T15:01:45Z", + "id": "v1_ChdXS0l4YWZXTk9xbk0xZThQczhEcmlROBIXV0tJeGFmV05PcW5NMWU4UHM4RHJpUTg", + "model": "gemini-3-flash-preview", + "object": "interaction", + "status": "completed", + "updated": "2025-12-04T15:01:45Z" + } + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "InteractionCreatedEvent": { + "description": "Server response confirming that a new interaction was created.", + "type": "object", + "properties": { + "interaction": { + "description": "Required. Partial interaction resource emitted when the stream is created.", + "$ref": "#/components/schemas/InteractionSseEventInteraction" + }, + "event_type": { + "const": "interaction.created" + }, + "event_id": { + "description": "The event_id token to be used to resume the interaction stream, from\nthis event.", + "type": "string" + } + }, + "required": [ + "event_type", + "interaction" + ], + "examples": [ + { + "interaction_created": { + "summary": "Interaction Created", + "value": { + "event_id": "evt_123", + "event_type": "interaction.created", + "interaction": { + "created": "2025-12-04T15:01:45Z", + "id": "v1_ChdXS0l4YWZXTk9xbk0xZThQczhEcmlROBIXV0tJeGFmV05PcW5NMWU4UHM4RHJpUTg", + "model": "gemini-3.6-flash", + "status": "in_progress", + "updated": "2025-12-04T15:01:45Z" + } + } + } + }, + { + "interaction_created": { + "summary": "Interaction Created", + "value": { + "event_id": "evt_123", + "event_type": "interaction.created", + "interaction": { + "id": "v1_ChdXS0l4YWZXTk9xbk0xZThQczhEcmlROBIXV0tJeGFmV05PcW5NMWU4UHM4RHJpUTg", + "model": "gemini-3-flash-preview", + "object": "interaction", + "status": "in_progress" + } + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "InteractionSseEvent": { + "type": "object", + "oneOf": [ + { + "$ref": "#/components/schemas/ErrorEvent" + }, + { + "$ref": "#/components/schemas/InteractionCompletedEvent" + }, + { + "$ref": "#/components/schemas/InteractionCreatedEvent" + }, + { + "$ref": "#/components/schemas/InteractionStatusUpdate" + }, + { + "$ref": "#/components/schemas/StepDelta" + }, + { + "$ref": "#/components/schemas/StepStart" + }, + { + "$ref": "#/components/schemas/StepStop" + } + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "InteractionSSEEvent" + }, + "InteractionSseEventInteraction": { + "description": "Partial interaction resource emitted by interaction lifecycle SSE events.\nStreaming lifecycle payloads may omit fields that are only available on\nfull non-streaming Interaction responses.\n", + "type": "object", + "properties": { + "id": { + "description": "Required. Output only. A unique identifier for the interaction completion.", + "type": "string", + "readOnly": true + }, + "status": { + "description": "Required. Output only. The status of the interaction.", + "type": "string", + "enum": [ + "in_progress", + "requires_action", + "completed", + "failed", + "cancelled", + "incomplete" + ], + "readOnly": true, + "x-google-enum-descriptions": [ + "The interaction is in progress.", + "The interaction requires action/input from the user.", + "The interaction is completed.", + "The interaction failed.", + "The interaction was cancelled.", + "The interaction is completed, but contains incomplete results (e.g. hitting max_tokens)." + ] + }, + "model": { + "description": "The model that will complete your prompt.", + "type": "string" + }, + "created": { + "description": "Output only. The time at which the response was created in ISO 8601 format.", + "type": "string", + "readOnly": true + }, + "updated": { + "description": "Output only. The time at which the response was last updated in ISO 8601 format.", + "type": "string", + "readOnly": true + }, + "usage": { + "description": "Output only. Statistics on the interaction request's token usage.", + "$ref": "#/components/schemas/Usage", + "readOnly": true + }, + "service_tier": { + "description": "The service tier for the interaction.", + "$ref": "#/components/schemas/ServiceTier" + }, + "steps": { + "description": "Output only. The steps that make up the interaction, if included in this event.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Step" + }, + "readOnly": true + }, + "agent": { + "description": "The agent to interact with.", + "type": "string" + }, + "continuation_token": { + "description": "Opaque token to resume a long decode.\nOutput: set when status is INCOMPLETE and decoding can be resumed.\nInput: pass the latest token back unchanged in CreateInteraction to\ncontinue decoding.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "object": { + "description": "Output only. The resource type.", + "type": "string", + "readOnly": true + } + }, + "required": [ + "id", + "status" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "InteractionSseStreamEvent": { + "type": "object", + "properties": { + "data": { + "$ref": "#/components/schemas/InteractionSseEvent" + } + }, + "required": [ + "data" + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "InteractionSSEStreamEvent" + }, + "InteractionStatusUpdate": { + "type": "object", + "properties": { + "interaction_id": { + "type": "string" + }, + "status": { + "type": "string", + "enum": [ + "in_progress", + "requires_action", + "completed", + "failed", + "cancelled", + "incomplete", + "budget_exceeded" + ], + "x-google-enum-deprecated": [ + false, + false, + false, + false, + false, + false, + true + ], + "x-google-enum-descriptions": [ + "The interaction is in progress.", + "The interaction requires action/input from the user.", + "The interaction is completed.", + "The interaction failed.", + "The interaction was cancelled.", + "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + ] + }, + "event_type": { + "const": "interaction.status_update" + }, + "continuation_token": { + "description": "An optional opaque continuation token used to resume decoding from the\nlatest checkpoint after a disconnected stream.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "usage": { + "description": "Model usage statistics up to this checkpoint.", + "$ref": "#/components/schemas/Usage" + }, + "event_id": { + "description": "The event_id token to be used to resume the interaction stream, from\nthis event.", + "type": "string" + } + }, + "required": [ + "event_type", + "interaction_id", + "status" + ], + "examples": [ + { + "interaction_status_update": { + "summary": "Interaction Status Update", + "value": { + "event_type": "interaction.status_update", + "interaction_id": "v1_ChdTMjQ0YWJ5TUF1TzcxZThQdjRpcnFRcxIXUzI0NGFieU1BdU83MWU4UHY0aXJxUXM", + "status": "in_progress" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "InteractionsInput": { + "description": "The input for the interaction.", + "oneOf": [ + { + "$ref": "#/components/schemas/Content" + }, + { + "title": "StepList", + "type": "array", + "items": { + "$ref": "#/components/schemas/Step" + } + }, + { + "title": "ContentList", + "type": "array", + "items": { + "$ref": "#/components/schemas/Content" + } + }, + { + "type": "string" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ListAgentsResponse": { + "type": "object", + "properties": { + "agents": { + "description": "The list of agents.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Agent" + } + }, + "next_page_token": { + "description": "A token to retrieve the next page of results.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ListCredentialsResponse": { + "type": "object", + "properties": { + "credentials": { + "type": "array", + "items": { + "$ref": "#/components/schemas/Credential" + } + }, + "next_page_token": { + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ListEnvironmentsResponse": { + "description": "Response for `ListEnvironments`.", + "type": "object", + "properties": { + "environments": { + "description": "Environments belonging to the provided project.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Environment" + } + }, + "next_page_token": { + "description": "Pagination token.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ListTriggerExecutionsResponse": { + "description": "Response message for TriggerService.ListTriggerExecutions.", + "type": "object", + "properties": { + "trigger_executions": { + "description": "The list of trigger executions.", + "type": "array", + "items": { + "$ref": "#/components/schemas/TriggerExecution" + } + }, + "next_page_token": { + "description": "A page token, received from a previous `ListTriggerExecutions` call.\nProvide this to retrieve the subsequent page.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ListTriggersResponse": { + "description": "Response message for TriggerService.ListTriggers.", + "type": "object", + "properties": { + "triggers": { + "description": "The list of triggers.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Trigger" + } + }, + "next_page_token": { + "description": "A page token, received from a previous `ListTriggers` call.\nProvide this to retrieve the subsequent page.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ListVoicesResponse": { + "description": "Response message for `VoicesService.ListVoices`.", + "type": "object", + "properties": { + "voices": { + "description": "The voices from the specified collection.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Voice" + } + }, + "next_page_token": { + "description": "A token that can be sent as `page_token` to retrieve the next page.\nIf empty, there are no subsequent pages.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ListWebhooksResponse": { + "description": "Response message for WebhookService.ListWebhooks.", + "type": "object", + "properties": { + "webhooks": { + "description": "The webhooks.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Webhook" + } + }, + "next_page_token": { + "description": "A token, which can be sent as `page_token` to retrieve the next page.\nIf this field is omitted, there are no subsequent pages.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "McpServer": { + "description": "A MCPServer is a server that can be called by the model to perform actions.", + "type": "object", + "properties": { + "type": { + "const": "mcp_server" + }, + "url": { + "description": "The full URL for the MCPServer endpoint.\nExample: \"https://api.example.com/mcp\"", + "type": "string" + }, + "name": { + "description": "The name of the MCPServer.", + "type": "string" + }, + "headers": { + "description": "Optional: Fields for authentication headers, timeouts, etc., if needed.", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "allowed_tools": { + "description": "The allowed tools.", + "type": "array", + "items": { + "$ref": "#/components/schemas/AllowedTools" + } + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-codeSamples": [ + { + "label": "mcp_server", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"tools\": [{\n \"type\": \"mcp_server\",\n \"name\": \"weather_service\",\n \"url\": \"https://gemini-api-demos.uc.r.appspot.com/mcp\"\n }],\n \"input\": \"Today is 12-05-2025, what is the temperature today in London?\"\n }'\n" + }, + { + "label": "mcp_server", + "lang": "python", + "source": "import os\n\nfrom google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n tools=[{\n \"type\": \"mcp_server\",\n \"name\": \"weather_service\",\n \"url\": \"https://gemini-api-demos.uc.r.appspot.com/mcp\",\n }],\n input=\"Today is 12-05-2025, what is the temperature today in London?\",\n)\nprint(response.output_text)\n" + }, + { + "label": "mcp_server", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n tools: [\n {\n type: 'mcp_server',\n name: 'weather_service',\n url: 'https://gemini-api-demos.uc.r.appspot.com/mcp',\n },\n ],\n input: 'Today is 12-05-2025, what is the temperature today in London?',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "mcp_server", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.MCPServer;\n\nClient client = new Client();\nMCPServer mcpTool =\n MCPServer.builder()\n .name(\"weather_service\")\n .url(\"https://gemini-api-demos.uc.r.appspot.com/mcp\")\n .build();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .tools(List.of(mcpTool))\n .input(\n InteractionsInput.of(\n \"Today is 12-05-2025, what is the temperature today in London?\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "mcp_server", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"os\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n mcpServer := gaos_interactions.MCPServer{\n Name: genai.Ptr(\"weather_service\"),\n URL: genai.Ptr(\"https://gemini-api-demos.uc.r.appspot.com/mcp\"),\n }\n tool := gaos_interactions.NewTool(mcpServer)\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Today is 12-05-2025, what is the temperature today in London?\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + } + ] + }, + "McpServerToolCallDelta": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "server_name": { + "type": "string" + }, + "arguments": { + "type": "object", + "additionalProperties": { + "description": "Properties of the object." + } + }, + "type": { + "const": "mcp_server_tool_call" + } + }, + "required": [ + "arguments", + "name", + "server_name", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "McpServerToolCallStep": { + "description": "MCPServer tool call step.", + "type": "object", + "properties": { + "name": { + "description": "Required. The name of the tool which was called.", + "type": "string" + }, + "server_name": { + "description": "Required. The name of the used MCP server.", + "type": "string" + }, + "arguments": { + "description": "Required. The JSON object of arguments for the function.", + "type": "object", + "additionalProperties": { + "description": "Properties of the object." + } + }, + "id": { + "description": "Required. A unique ID for this specific tool call.", + "type": "string" + }, + "type": { + "const": "mcp_server_tool_call" + } + }, + "required": [ + "arguments", + "id", + "name", + "server_name", + "type" + ], + "examples": [ + { + "mcp_server_tool_call": { + "summary": "McpServerToolCallStep", + "value": { + "name": "calculate_tax", + "type": "mcp_server_tool_call", + "arguments": { + "income": 120000, + "state": "CA" + }, + "id": "mcp_call_29012", + "server_name": "financial_mcp_server" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "McpServerToolResultDelta": { + "type": "object", + "properties": { + "result": { + "oneOf": [ + { + "type": "array", + "items": { + "oneOf": [ + { + "$ref": "#/components/schemas/ImageContent" + }, + { + "$ref": "#/components/schemas/TextContent" + } + ] + } + }, + { + "type": "object" + }, + { + "type": "string" + } + ] + }, + "type": { + "const": "mcp_server_tool_result" + }, + "name": { + "type": "string" + }, + "server_name": { + "type": "string" + } + }, + "required": [ + "result", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "McpServerToolResultStep": { + "description": "MCPServer tool result step.", + "type": "object", + "properties": { + "result": { + "description": "Required. The output from the MCP server call. Can be simple text or rich content.", + "oneOf": [ + { + "type": "array", + "items": { + "oneOf": [ + { + "$ref": "#/components/schemas/ImageContent" + }, + { + "$ref": "#/components/schemas/TextContent" + } + ] + } + }, + { + "type": "object" + }, + { + "type": "string" + } + ] + }, + "call_id": { + "description": "Required. ID to match the ID from the function call block.", + "type": "string" + }, + "type": { + "const": "mcp_server_tool_result" + }, + "name": { + "description": "Name of the tool which is called for this specific tool call.", + "type": "string" + }, + "server_name": { + "description": "The name of the used MCP server.", + "type": "string" + } + }, + "required": [ + "call_id", + "result", + "type" + ], + "examples": [ + { + "mcp_server_tool_result": { + "summary": "McpServerToolResultStep", + "value": { + "name": "calculate_tax", + "type": "mcp_server_tool_result", + "call_id": "mcp_call_29012", + "result": { + "tax_due": 32400 + }, + "server_name": "financial_mcp_server" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "MediaProcessing": { + "oneOf": [ + { + "$ref": "#/components/schemas/StaticMediaProcessing" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "MediaResolution": { + "type": "string", + "enum": [ + "low", + "medium", + "high", + "ultra_high" + ], + "x-google-enum-descriptions": [ + "Low resolution.", + "Medium resolution.", + "High resolution.", + "Ultra high resolution." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ModalityTokens": { + "description": "The token count for a single response modality.", + "type": "object", + "properties": { + "modality": { + "description": "The modality associated with the token count.", + "$ref": "#/components/schemas/ResponseModality" + }, + "tokens": { + "description": "Number of tokens for the modality.", + "type": "integer", + "format": "int32" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ModelOption": { + "title": "Model", + "description": "The model that will complete your prompt.\\n\\nSee [models](https://ai.google.dev/gemini-api/docs/models) for additional details.", + "type": "string", + "enum": [ + "gemma-4-26b-a4b-it", + "gemma-4-31b-it", + "gemini-flash-latest", + "gemini-flash-lite-latest", + "gemini-pro-latest", + "gemini-2.5-flash", + "gemini-2.5-pro", + "gemini-2.5-flash-lite", + "gemini-3.5-flash-lite", + "gemini-2.5-flash-image", + "gemini-3-flash-preview", + "gemini-3.1-pro-preview", + "gemini-3.1-pro-preview-customtools", + "gemini-3.1-flash-lite", + "gemini-3.1-flash-lite-image", + "gemini-3-pro-image", + "nano-banana-pro-preview", + "gemini-3.1-flash-image", + "gemini-nano-banana-2.1", + "gemini-3.1-flash-tts-preview", + "gemini-3.5-flash", + "gemini-3.6-flash", + "gemini-3.7-flash", + "gemini-3.8-flash", + "gemini-3.8-flash-tts", + "gemini-3.8-flash-lite-tts", + "lyria-3-clip-preview", + "lyria-3-pro-preview", + "gemini-robotics-er-1.6-preview", + "gemini-robotics-er-2-preview", + "lyria-3.5", + "gemini-omni-1.1-flash", + "gemini-omni-flash-preview" + ], + "x-speakeasy-enum-descriptions": [ + "Gemma 4 26B A4B IT", + "Gemma 4 31B IT", + "Latest release of Gemini Flash", + "Latest release of Gemini Flash-Lite", + "Latest release of Gemini Pro", + "Our first hybrid reasoning model which supports a 1M token context window and has thinking budgets.", + "Our state-of-the-art multipurpose model, which excels at coding and complex reasoning tasks.", + "Our smallest and most cost effective model, built for at scale usage.", + "Gemini 3.5 Flash-Lite - Our fastest, most cost-effective 3.5 model for high-throughput execution.", + "Our native image generation model, optimized for speed, flexibility, and contextual understanding. Text input and output is priced the same as 2.5 Flash.", + "Our most intelligent model built for speed, combining frontier intelligence with superior search and grounding.", + "Our latest SOTA reasoning model with unprecedented depth and nuance, and powerful multimodal understanding and coding capabilities.", + "Gemini 3.1 Pro Preview optimized for custom tool usage", + "Our most cost-efficient model, optimized for high-volume agentic tasks, translation, and simple data processing.", + "Gemini 3.1 Flash Lite Image.", + "Gemini 3 Pro Image", + "Gemini 3 Pro Image Preview", + "Gemini 3.1 Flash Image.", + "Gemini Nano Banana 2.1.", + "Gemini 3.1 Flash TTS: Powerful, low-latency speech generation. Enjoy natural outputs, steerable prompts, and new expressive audio tags for precise narration control.", + "Gemini 3.5 Flash - Our earlier Flash model, built for speed and foundational performance across routine, high-throughput workloads.", + "Gemini 3.6 Flash - Our previous generation Flash model, balancing speed and multimodal capabilities across general agentic and everyday tasks.", + "Gemini 3.7 Flash - Our high-speed, efficient Flash model built for everyday coding, agentic tool use, and reliable multi-step execution.", + "Gemini 3.8 Flash - Our most intelligent Flash model, engineered for long-horizon software engineering, autonomous agents, and complex enterprise workflows.", + "Gemini 3.8 Flash TTS - Flagship TTS model for Voice Design and dual-speaker screenplay control. Prompt custom vocal personas, direct line-by-line delivery, and add vocal bursts.", + "Gemini 3.8 Flash Lite TTS - High-speed and cost-efficient, ideal for rapid dubbing, media localization, and high-throughput voice agents. Direct replacement for gemini-3.1-flash-tts-preview.", + "Our low-latency, music generation model optimized for high-fidelity audio clips and precise rhythmic control.", + "Our advanced, full-song generative model with deep compositional understanding, optimized for precise structural control and complex transitions across diverse musical styles.", + "Gemini Robotics-ER 1.6 Preview", + "Gemini Robotics Embodied Reasoning 2 Preview", + "Our flagship music generation model, optimized for full-length songs with complex structural coherence.", + "Our high-performance multimodal model designed for fast, conversational video generation, editing, and cinematic control.", + "Preview release of our multimodal model for conversational video generation, editing, and cinematic control." + ], + "x-speakeasy-enum-format": "union", + "x-speakeasy-exports": [ + { + "name": "model", + "group": "interactions" + }, + { + "name": "model_param", + "group": "interactions" + } + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "Model", + "x-speakeasy-unknown-values": "allow" + }, + "ModelOutputStep": { + "description": "Output generated by the model.", + "type": "object", + "properties": { + "type": { + "const": "model_output" + }, + "content": { + "type": "array", + "items": { + "$ref": "#/components/schemas/Content" + } + } + }, + "required": [ + "type" + ], + "examples": [ + { + "model_output": { + "summary": "ModelOutputStep", + "value": { + "type": "model_output", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris." + } + ] + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "OAuth2Config": { + "description": "Configuration for OAuth2 credentials with automatic token refresh.", + "type": "object", + "properties": { + "client_id": { + "description": "Required. OAuth2 client ID.", + "type": "string" + }, + "client_secret": { + "description": "Required. Input only. OAuth2 client secret. Write-only; never returned in responses.", + "type": "string", + "writeOnly": true + }, + "refresh_token": { + "description": "Required. Input only. OAuth2 refresh token. Write-only; never returned in responses.", + "type": "string", + "writeOnly": true + }, + "token_url": { + "description": "Required. OAuth2 token endpoint URL for refreshing access tokens.", + "type": "string" + }, + "id": { + "type": "string" + }, + "type": { + "const": "oauth2" + }, + "scopes": { + "description": "Optional. List of OAuth2 scopes.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "client_id", + "client_secret", + "id", + "refresh_token", + "token_url", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "OAuth2UpdateConfig": { + "description": "Configuration for updating OAuth2 credentials.", + "type": "object", + "properties": { + "type": { + "const": "oauth2" + }, + "client_id": { + "description": "Optional. OAuth2 client ID.", + "type": "string" + }, + "client_secret": { + "description": "Optional. Input only. OAuth2 client secret. Write-only; never returned in responses.", + "type": "string", + "writeOnly": true + }, + "refresh_token": { + "description": "Optional. Input only. OAuth2 refresh token. Write-only; never returned in responses.", + "type": "string", + "writeOnly": true + }, + "token_url": { + "description": "Optional. OAuth2 token endpoint URL for refreshing access tokens.", + "type": "string" + }, + "scopes": { + "description": "Optional. List of OAuth2 scopes.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "PingWebhookRequest": { + "description": "Request message for WebhookService.PingWebhook.", + "type": "object", + "properties": { + "id": { + "description": "Required. The ID of the webhook to ping.\nFormat: `{webhook_id}`", + "type": "string" + } + }, + "required": [ + "id" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "PingWebhookResponse": { + "description": "Response message for WebhookService.PingWebhook.", + "type": "object", + "x-speakeasy-model-namespace": "interactions" + }, + "Pitch": { + "type": "string", + "enum": [ + "low", + "medium", + "high" + ], + "x-google-enum-descriptions": [ + "Lower pitch voice.", + "Medium pitch voice.", + "Higher pitch voice." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "PlaceCitation": { + "description": "A place citation annotation.", + "type": "object", + "properties": { + "type": { + "const": "place_citation" + }, + "place_id": { + "description": "The ID of the place, in `places/{place_id}` format.", + "type": "string" + }, + "name": { + "description": "Title of the place.", + "type": "string" + }, + "url": { + "description": "URI reference of the place.", + "type": "string" + }, + "review_snippets": { + "description": "Snippets of reviews that are used to generate answers about the\nfeatures of a given place in Google Maps.", + "type": "array", + "items": { + "$ref": "#/components/schemas/ReviewSnippet" + } + }, + "start_index": { + "description": "Start of segment of the response that is attributed to this source.\n\nIndex indicates the start of the segment, measured in bytes.", + "type": "integer", + "format": "int32" + }, + "end_index": { + "description": "End of the attributed segment, exclusive.", + "type": "integer", + "format": "int32" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Places": { + "type": "object", + "properties": { + "place_id": { + "description": "The ID of the place, in `places/{place_id}` format.", + "type": "string" + }, + "name": { + "description": "Title of the place.", + "type": "string" + }, + "url": { + "description": "URI reference of the place.", + "type": "string" + }, + "review_snippets": { + "description": "Snippets of reviews that are used to generate answers about the\nfeatures of a given place in Google Maps.", + "type": "array", + "items": { + "$ref": "#/components/schemas/ReviewSnippet" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "ProcessingCallDelta": { + "description": "Streaming delta for a server-initiated media processing step.", + "type": "object", + "properties": { + "type": { + "const": "processing_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ProcessingCallStep": { + "description": "A server-initiated processing step for media analysis (e.g. video\nunderstanding).", + "type": "object", + "properties": { + "id": { + "description": "Required. A unique ID for this specific tool call.", + "type": "string" + }, + "type": { + "const": "processing_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "id", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ProcessingResultDelta": { + "description": "Streaming delta for the result of a server-initiated media processing step.", + "type": "object", + "properties": { + "type": { + "const": "processing_result" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ProcessingResultStep": { + "description": "The result of a server-initiated media processing step.", + "type": "object", + "properties": { + "call_id": { + "description": "Required. ID to match the ID from the function call block.", + "type": "string" + }, + "type": { + "const": "processing_result" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "call_id", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "PromptedVoice": { + "description": "Parameters for prompted voice generation.\nRequired in `CreateVoice` when `type` is `\"prompted\"`. Returned in\n`CreateVoice`, `GetVoice`, and `ListVoices` responses for prompted voices.", + "type": "object", + "properties": { + "input": { + "description": "Required. The natural-language prompt describing the desired voice, e.g.\n\"A deep, booming male voice of a massive evil ogre in his middle years.\"", + "type": "string" + }, + "base_voice": { + "description": "Optional. The name or ID of an existing base voice to edit or remix.\nUse the voice name for prebuilt voices (e.g., `Puck`) or the voice ID for\npreviously created voices (e.g., `voice_abc123`). This can be any\nprompted or replicated voice that the user has access to.\nWhen present, `input` describes relative alterations to make to the base\nvoice.", + "type": "string" + } + }, + "required": [ + "input" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ReplicatedVoice": { + "description": "Input only. Parameters for replicated voice generation.\nRequired on input when `type` is `\"replicated\"`. Not returned in\nresponses.", + "type": "object", + "properties": { + "source_audio": { + "description": "Required. The reference audio sample of the speaker's voice to replicate.\nMust be an unmodified recording of an adult human voice.", + "$ref": "#/components/schemas/AudioData" + }, + "consent_audio": { + "description": "Required. Audio recording of the same speaker reading the required\nconsent phrase, confirming authorization to replicate the voice.", + "$ref": "#/components/schemas/AudioData" + } + }, + "required": [ + "consent_audio", + "source_audio" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ResponseFormat": { + "type": "object", + "oneOf": [ + { + "$ref": "#/components/schemas/AudioResponseFormat" + }, + { + "$ref": "#/components/schemas/ImageResponseFormat" + }, + { + "$ref": "#/components/schemas/TextResponseFormat" + }, + { + "$ref": "#/components/schemas/VideoResponseFormat" + }, + { + "type": "object", + "additionalProperties": true + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ResponseModality": { + "type": "string", + "enum": [ + "text", + "image", + "audio", + "video", + "document" + ], + "x-google-enum-descriptions": [ + "Indicates the model should return text.", + "Indicates the model should return images.", + "Indicates the model should return audio.", + "Indicates the model should return video.", + "Indicates the model should return documents." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ReviewSnippet": { + "description": "Encapsulates a snippet of a user review that answers a question about\nthe features of a specific place in Google Maps.", + "type": "object", + "properties": { + "title": { + "description": "Title of the review.", + "type": "string" + }, + "url": { + "description": "A link that corresponds to the user review on Google Maps.", + "type": "string" + }, + "review_id": { + "description": "The ID of the review snippet.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "RotateSigningSecretRequest": { + "description": "Request message for WebhookService.RotateSigningSecret.", + "type": "object", + "properties": { + "id": { + "description": "Required. The ID of the webhook for which to generate a signing secret.\nFormat: `{webhook_id}`", + "type": "string" + }, + "revocation_behavior": { + "description": "Optional. The revocation behavior for previous signing secrets.", + "type": "string", + "enum": [ + "revoke_previous_secrets_after_h24", + "revoke_previous_secrets_immediately" + ], + "x-google-enum-descriptions": [ + "Generate a new signing secret and revoke all previous secrets after 24\nhours. Default and safest option for migrations.", + "Revoke all previous secrets immediately. Use with caution as this can\ninterrupt ongoing notifications." + ] + } + }, + "required": [ + "id" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "RotateSigningSecretResponse": { + "description": "Response message for WebhookService.RotateSigningSecret.", + "type": "object", + "properties": { + "secret": { + "description": "Output only. The newly generated signing secret.", + "type": "string", + "readOnly": true + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "RunTriggerRequest": { + "description": "Request message for TriggerService.RunTrigger.", + "type": "object", + "properties": { + "id": { + "description": "Required. The ID of the trigger to run immediately.", + "type": "string" + } + }, + "required": [ + "id" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "SafetySetting": { + "description": "A safety setting that affects the safety-blocking behavior.\n\nA SafetySetting consists of a\nharm category and a\nthreshold for that\ncategory.", + "type": "object", + "properties": { + "type": { + "description": "Required. The type of harm category to be blocked.", + "$ref": "#/components/schemas/HarmCategory" + }, + "threshold": { + "description": "Required. The threshold for blocking content. If the harm probability\nexceeds this threshold, the content will be blocked.", + "type": "string", + "enum": [ + "block_low_and_above", + "block_medium_and_above", + "block_only_high", + "off" + ], + "x-google-enum-descriptions": [ + "Block content with a low harm probability or higher.", + "Block content with a medium harm probability or higher.", + "Block content with a high harm probability.", + "Turn off the safety filter entirely." + ] + }, + "method": { + "description": "Optional. The method for blocking content. If not specified, the default\nbehavior is to use the probability score.", + "type": "string", + "enum": [ + "severity", + "probability" + ], + "x-google-enum-descriptions": [ + "The harm block method uses both probability and severity scores.", + "The harm block method uses the probability score." + ] + } + }, + "required": [ + "threshold", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ServiceTier": { + "type": "string", + "enum": [ + "flex", + "standard", + "priority" + ], + "x-google-enum-descriptions": [ + "Flex service tier.", + "Standard service tier.", + "Priority service tier." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "SigningSecret": { + "description": "Represents a signing secret used to verify webhook payloads.", + "type": "object", + "properties": { + "truncated_secret": { + "description": "Output only. The truncated version of the signing secret.", + "type": "string", + "readOnly": true + }, + "expire_time": { + "description": "Output only. The expiration date of the signing secret.", + "type": "string", + "format": "date-time", + "readOnly": true + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "SmartTranscriptionMode": { + "description": "Configuration for smart transcription mode.", + "type": "object", + "properties": { + "type": { + "const": "smart" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Source": { + "description": "A source to be mounted into the environment.", + "type": "object", + "properties": { + "type": { + "type": "string", + "enum": [ + "gcs", + "inline", + "repository" + ], + "x-google-enum-descriptions": [ + "A Cloud Storage bucket.", + "Inline content.", + "A generic repository. The protocol prefix in the source URL\nidentifies the provider (e.g., github://, gcs://)." + ] + }, + "source": { + "description": "The source of the environment.\nFor Cloud Storage, this is the Cloud Storage path.\nFor GitHub, this is the GitHub path.", + "type": "string" + }, + "target": { + "description": "Where the source should appear in the environment.", + "type": "string" + }, + "content": { + "description": "The inline content if `type` is `INLINE`.", + "type": "string" + }, + "encoding": { + "description": "Optional encoding for inline content (e.g. `base64`).", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "SpeakerConfig": { + "description": "Configuration for multi-speaker and speech generation.", + "type": "object", + "properties": { + "speakers": { + "description": "Individual speaker configurations.", + "type": "array", + "items": { + "$ref": "#/components/schemas/SpeechConfig" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "SpeechAnnotation": { + "description": "Speech annotation for text content.", + "type": "object", + "properties": { + "type": { + "const": "speech_metadata" + }, + "speaker": { + "description": "The speaker to associate with this turn.", + "type": "string" + }, + "style": { + "description": "Style instruction for the speech synthesis.", + "type": "string" + }, + "start_index": { + "description": "Start of segment of the response that is attributed to this source.\n\nIndex indicates the start of the segment, measured in bytes.", + "type": "integer", + "format": "int32" + }, + "end_index": { + "description": "End of the attributed segment, exclusive.", + "type": "integer", + "format": "int32" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "SpeechConfig": { + "description": "The configuration for speech interaction.", + "type": "object", + "properties": { + "voice": { + "description": "The voice of the speaker.", + "type": "string" + }, + "language": { + "description": "The language of the speech.", + "type": "string" + }, + "speaker": { + "description": "The speaker's name, it should match the speaker name given in the prompt.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "StaticMediaProcessing": { + "type": "object", + "properties": { + "type": { + "const": "static" + }, + "start_offset": { + "description": "Optional. Segment start time. Specified as a decimal number of seconds followed\nby an 's' suffix, e.g., \"10.5s\". Must be non-negative.", + "type": "string", + "format": "google-duration" + }, + "end_offset": { + "description": "Optional. Segment end time. Specified as a decimal number of seconds followed\nby an 's' suffix, e.g., \"30s\". Must be non-negative and greater than\n`start_offset` if `start_offset` is set.", + "type": "string", + "format": "google-duration" + }, + "fps": { + "description": "Optional. Video frame-rate sampling density.", + "type": "number", + "format": "double" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Step": { + "description": "A step in the interaction.", + "type": "object", + "oneOf": [ + { + "$ref": "#/components/schemas/CodeExecutionCallStep" + }, + { + "$ref": "#/components/schemas/CodeExecutionResultStep" + }, + { + "$ref": "#/components/schemas/FileSearchCallStep" + }, + { + "$ref": "#/components/schemas/FileSearchResultStep" + }, + { + "$ref": "#/components/schemas/FunctionCallStep" + }, + { + "$ref": "#/components/schemas/FunctionResultStep" + }, + { + "$ref": "#/components/schemas/GoogleMapsCallStep" + }, + { + "$ref": "#/components/schemas/GoogleMapsResultStep" + }, + { + "$ref": "#/components/schemas/GoogleSearchCallStep" + }, + { + "$ref": "#/components/schemas/GoogleSearchResultStep" + }, + { + "$ref": "#/components/schemas/McpServerToolCallStep" + }, + { + "$ref": "#/components/schemas/McpServerToolResultStep" + }, + { + "$ref": "#/components/schemas/ModelOutputStep" + }, + { + "$ref": "#/components/schemas/ProcessingCallStep" + }, + { + "$ref": "#/components/schemas/ProcessingResultStep" + }, + { + "$ref": "#/components/schemas/ThoughtStep" + }, + { + "$ref": "#/components/schemas/UrlContextCallStep" + }, + { + "$ref": "#/components/schemas/UrlContextResultStep" + }, + { + "$ref": "#/components/schemas/UserInputStep" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "StepDelta": { + "type": "object", + "properties": { + "index": { + "type": "integer", + "format": "int32" + }, + "delta": { + "$ref": "#/components/schemas/StepDeltaData" + }, + "event_type": { + "const": "step.delta" + }, + "event_id": { + "description": "The event_id token to be used to resume the interaction stream, from\nthis event.", + "type": "string" + }, + "metadata": { + "$ref": "#/components/schemas/StepDeltaMetadata" + } + }, + "required": [ + "delta", + "event_type", + "index" + ], + "examples": [ + { + "step_delta": { + "summary": "Step Delta", + "value": { + "delta": { + "type": "text", + "text": "Hello" + }, + "event_type": "step.delta", + "index": 0 + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "StepDeltaData": { + "type": "object", + "oneOf": [ + { + "$ref": "#/components/schemas/ArgumentsDelta" + }, + { + "$ref": "#/components/schemas/AudioDelta" + }, + { + "$ref": "#/components/schemas/CodeExecutionCallDelta" + }, + { + "$ref": "#/components/schemas/CodeExecutionResultDelta" + }, + { + "$ref": "#/components/schemas/DocumentDelta" + }, + { + "$ref": "#/components/schemas/FileSearchCallDelta" + }, + { + "$ref": "#/components/schemas/FileSearchResultDelta" + }, + { + "$ref": "#/components/schemas/FunctionResultDelta" + }, + { + "$ref": "#/components/schemas/GoogleMapsCallDelta" + }, + { + "$ref": "#/components/schemas/GoogleMapsResultDelta" + }, + { + "$ref": "#/components/schemas/GoogleSearchCallDelta" + }, + { + "$ref": "#/components/schemas/GoogleSearchResultDelta" + }, + { + "$ref": "#/components/schemas/ImageDelta" + }, + { + "$ref": "#/components/schemas/McpServerToolCallDelta" + }, + { + "$ref": "#/components/schemas/McpServerToolResultDelta" + }, + { + "$ref": "#/components/schemas/ProcessingCallDelta" + }, + { + "$ref": "#/components/schemas/ProcessingResultDelta" + }, + { + "$ref": "#/components/schemas/TextAnnotationDelta" + }, + { + "$ref": "#/components/schemas/TextDelta" + }, + { + "$ref": "#/components/schemas/ThoughtSignatureDelta" + }, + { + "$ref": "#/components/schemas/ThoughtSummaryDelta" + }, + { + "$ref": "#/components/schemas/UrlContextCallDelta" + }, + { + "$ref": "#/components/schemas/UrlContextResultDelta" + }, + { + "$ref": "#/components/schemas/VideoDelta" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "StepDeltaMetadata": { + "description": "Optional metadata accompanying ANY streamed event.", + "type": "object", + "properties": { + "total_usage": { + "description": "Statistics on the interaction request's token usage.", + "$ref": "#/components/schemas/Usage" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "StepStart": { + "type": "object", + "properties": { + "index": { + "type": "integer", + "format": "int32" + }, + "step": { + "$ref": "#/components/schemas/Step" + }, + "event_type": { + "const": "step.start" + }, + "event_id": { + "description": "The event_id token to be used to resume the interaction stream, from\nthis event.", + "type": "string" + } + }, + "required": [ + "event_type", + "index", + "step" + ], + "examples": [ + { + "step_start": { + "summary": "Step Start", + "value": { + "event_type": "step.start", + "index": 0, + "step": { + "type": "model_output" + } + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "StepStop": { + "type": "object", + "properties": { + "index": { + "type": "integer", + "format": "int32" + }, + "event_type": { + "const": "step.stop" + }, + "usage": { + "description": "Cumulative model usage stats from the start of the session.", + "$ref": "#/components/schemas/Usage" + }, + "step_usage": { + "description": "Model usage stats for this specific step.", + "$ref": "#/components/schemas/Usage" + }, + "event_id": { + "description": "The event_id token to be used to resume the interaction stream, from\nthis event.", + "type": "string" + } + }, + "required": [ + "event_type", + "index" + ], + "examples": [ + { + "step_stop": { + "summary": "Step Stop", + "value": { + "event_type": "step.stop", + "index": 0 + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "TextAnnotationDelta": { + "type": "object", + "properties": { + "type": { + "const": "text_annotation_delta" + }, + "annotations": { + "description": "Citation information for model-generated content.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Annotation" + } + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "TextContent": { + "description": "A text content block.", + "type": "object", + "properties": { + "text": { + "description": "Required. The text content.", + "type": "string" + }, + "type": { + "const": "text" + }, + "annotations": { + "description": "Citation information for model-generated content.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Annotation" + } + } + }, + "required": [ + "text", + "type" + ], + "examples": [ + { + "text": { + "summary": "Text", + "value": { + "type": "text", + "text": "Hello, how are you?" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "TextDelta": { + "type": "object", + "properties": { + "text": { + "type": "string" + }, + "type": { + "const": "text" + } + }, + "required": [ + "text", + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "TextResponseFormat": { + "description": "Configuration for text output format.", + "type": "object", + "properties": { + "type": { + "const": "text" + }, + "mime_type": { + "description": "The MIME type of the text output.", + "type": "string", + "enum": [ + "application/json", + "text/plain" + ], + "x-google-enum-descriptions": [ + "JSON output format.", + "Plain text output format." + ] + }, + "schema": { + "description": "The JSON schema that the output should conform to. Only applicable when\nmime_type is application/json.\n\nSupports a subset of JSON Schema. `minLength`, `maxLength`,\n`minProperties`, `maxProperties`, `uniqueItems`, `contains`,\n`minContains`, and `maxContains` are ignored and not enforced; validate\nthese in your application. `minItems` and `maxItems` are enforced. Schemas\nthat are too complex to enforce (for example, large `maxItems` on nested\narrays) return a 400 INVALID_ARGUMENT error.", + "type": "object", + "additionalProperties": { + "description": "Properties of the object." + } + } + }, + "required": [ + "type" + ], + "examples": [ + { + "text_response_format": { + "summary": "Text Output (JSON Schema)", + "value": { + "type": "text", + "mime_type": "application/json", + "schema": { + "type": "object", + "properties": { + "ingredients": { + "type": "array", + "items": { + "type": "string" + } + }, + "recipe_name": { + "type": "string" + } + }, + "required": [ + "ingredients", + "recipe_name" + ] + } + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ThinkingLevel": { + "type": "string", + "enum": [ + "minimal", + "low", + "medium", + "high" + ], + "x-google-enum-descriptions": [ + "Little to no thinking.", + "Low thinking level.", + "Medium thinking level.", + "High thinking level." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ThinkingSummaries": { + "type": "string", + "enum": [ + "auto", + "none" + ], + "x-google-enum-descriptions": [ + "Auto thinking summaries.", + "No thinking summaries." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ThoughtSignatureDelta": { + "type": "object", + "properties": { + "type": { + "const": "thought_signature" + }, + "signature": { + "description": "Signature to match the backend source to be part of the generation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ThoughtStep": { + "description": "A thought step.", + "type": "object", + "properties": { + "type": { + "const": "thought" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "summary": { + "description": "A summary of the thought.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Content" + } + } + }, + "required": [ + "type" + ], + "examples": [ + { + "thought": { + "summary": "ThoughtStep", + "value": { + "type": "thought", + "signature": "thought_sig_abcd1234", + "summary": [ + { + "type": "text", + "text": "The model is searching Google for the capital of France." + } + ] + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ThoughtSummaryContent": { + "type": "object", + "oneOf": [ + { + "$ref": "#/components/schemas/ImageContent" + }, + { + "$ref": "#/components/schemas/TextContent" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ThoughtSummaryDelta": { + "type": "object", + "properties": { + "type": { + "const": "thought_summary" + }, + "content": { + "description": "A new summary item to be added to the thought.", + "$ref": "#/components/schemas/Content" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Tool": { + "description": "A tool that can be used by the model.", + "type": "object", + "oneOf": [ + { + "$ref": "#/components/schemas/CodeExecution" + }, + { + "$ref": "#/components/schemas/ComputerUse" + }, + { + "$ref": "#/components/schemas/FileSearch" + }, + { + "$ref": "#/components/schemas/Function" + }, + { + "$ref": "#/components/schemas/GoogleMaps" + }, + { + "$ref": "#/components/schemas/GoogleSearch" + }, + { + "$ref": "#/components/schemas/McpServer" + }, + { + "$ref": "#/components/schemas/UrlContext" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "ToolChoiceConfig": { + "description": "The tool choice configuration containing allowed tools.", + "type": "object", + "properties": { + "allowed_tools": { + "description": "The allowed tools.", + "$ref": "#/components/schemas/AllowedTools" + } + }, + "example": { + "allowed_tools": { + "mode": "any", + "tools": [ + "my_tool" + ] + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "TranscriptionConfig": { + "description": "Configuration for speech recognition (transcription).", + "type": "object", + "properties": { + "mode": { + "description": "Discriminated transcription mode options or enum.", + "oneOf": [ + { + "$ref": "#/components/schemas/TranscriptionMode" + }, + { + "type": "string", + "enum": [ + "verbatim", + "smart" + ], + "x-google-enum-descriptions": [ + "Verbatim transcription mode.", + "Smart transcription mode." + ] + } + ] + }, + "language_codes": { + "description": "Optional. BCP-47 language codes providing hints about the languages present in the\naudio. If omitted or empty, defaults to automatic language detection.", + "type": "array", + "items": { + "type": "string" + } + }, + "custom_vocabulary": { + "description": "Optional. A list of custom vocabulary phrases to bias the speech recognition model\ntoward recognizing specific terms.", + "type": "array", + "items": { + "type": "string" + } + }, + "timestamp_granularities": { + "description": "Optional. The granularity of timestamps to include in the transcription output.\nSupported values: \"word\". If empty, no timestamps are generated.", + "type": "array", + "deprecated": true, + "items": { + "type": "string" + } + }, + "diarization_mode": { + "description": "Optional. Configures speaker diarization. Supported values: \"speaker\".", + "type": "string", + "deprecated": true + }, + "adaptation_phrases": { + "description": "Optional. A list of phrases to bias the ASR model towards.", + "type": "array", + "deprecated": true, + "items": { + "type": "string" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "TranscriptionMode": { + "description": "Configuration for transcription mode.", + "oneOf": [ + { + "$ref": "#/components/schemas/SmartTranscriptionMode" + }, + { + "$ref": "#/components/schemas/VerbatimTranscriptionMode" + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Trigger": { + "description": "A trigger configuration that is scheduled to run an agent.", + "type": "object", + "properties": { + "schedule": { + "description": "Required. The cron schedule on which the trigger should run.\nStandard cron format.", + "type": "string" + }, + "time_zone": { + "description": "Required. Time zone in which the schedule should be interpreted.", + "type": "string" + }, + "interaction": { + "description": "Required. The interaction request template to be executed.", + "$ref": "#/components/schemas/CreateAgentInteractionParams" + }, + "id": { + "description": "Required. Output only. Identifier. The ID of the trigger.", + "type": "string", + "readOnly": true, + "x-google-identifier": true + }, + "display_name": { + "description": "Optional. The display name of the trigger.", + "type": "string" + }, + "environment_id": { + "description": "Optional. The environment ID for the trigger execution.", + "type": "string" + }, + "max_consecutive_failures": { + "description": "Optional. The maximum number of consecutive failures allowed before\nthe trigger is automatically paused (status becomes ERROR).", + "type": "integer", + "format": "int32" + }, + "execution_timeout_seconds": { + "description": "Optional. The execution timeout for the triggered interaction.", + "type": "integer", + "format": "int32" + }, + "consecutive_failure_count": { + "description": "Output only. The number of consecutive failures that have occurred\nsince the last successful execution.", + "type": "integer", + "format": "int32", + "readOnly": true + }, + "create_time": { + "description": "Output only. The time when the trigger was created.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "update_time": { + "description": "Output only. The time when the trigger was last updated.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "last_run_time": { + "description": "Output only. The time when the trigger was last run.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "next_run_time": { + "description": "Output only. The time when the trigger is scheduled to run next.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "last_pause_time": { + "description": "Output only. The time when the trigger was last paused.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "last_resume_time": { + "description": "Output only. The time when the trigger was last resumed.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "previous_interaction_id": { + "description": "Output only. The ID of the last interaction created by this trigger.", + "type": "string", + "readOnly": true + }, + "status": { + "description": "Output only. The current status of the trigger.", + "type": "string", + "enum": [ + "active", + "paused", + "error" + ], + "readOnly": true, + "x-google-enum-descriptions": [ + "The trigger is active and will fire on schedule.", + "The trigger is paused and will not fire.", + "The trigger has entered an error state due to consecutive failures." + ] + } + }, + "required": [ + "id", + "interaction", + "schedule", + "time_zone" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "TriggerExecution": { + "description": "An execution instance of a trigger.", + "type": "object", + "properties": { + "id": { + "description": "Required. Output only. Identifier. The ID of the trigger execution.", + "type": "string", + "readOnly": true, + "x-google-identifier": true + }, + "trigger_id": { + "description": "Required. Output only. Identifier. The ID of the trigger that created this execution.", + "type": "string", + "readOnly": true, + "x-google-identifier": true + }, + "scheduled_time": { + "description": "Output only. The time when the execution was scheduled to run.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "start_time": { + "description": "Output only. The time when the execution started.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "end_time": { + "description": "Output only. The time when the execution finished.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "status": { + "description": "Output only. The status of the execution.", + "type": "string", + "enum": [ + "in_progress", + "completed", + "failed", + "skipped", + "timed_out" + ], + "readOnly": true, + "x-google-enum-descriptions": [ + "The execution is currently in progress.", + "The execution completed successfully.", + "The execution failed.", + "The execution was skipped (e.g., previous execution still running).", + "The execution timed out." + ] + }, + "interaction_id": { + "description": "Output only. The ID of the interaction created by this execution, if any.", + "type": "string", + "readOnly": true + }, + "environment_id": { + "description": "Output only. The environment ID used for the execution.", + "type": "string", + "readOnly": true + }, + "error": { + "description": "Output only. The error message if the execution failed.", + "type": "string", + "readOnly": true + } + }, + "required": [ + "id", + "trigger_id" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "TriggerUpdate": { + "description": "Represents the fields of a Trigger that can be updated.", + "type": "object", + "properties": { + "status": { + "description": "Optional. The status of the trigger.", + "type": "string", + "enum": [ + "active", + "paused", + "error" + ], + "x-google-enum-descriptions": [ + "The trigger is active and will fire on schedule.", + "The trigger is paused and will not fire.", + "The trigger has entered an error state due to consecutive failures." + ] + }, + "display_name": { + "description": "Optional. The display name of the trigger.", + "type": "string" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "UpdateTriggerRequest": { + "description": "Request message for TriggerService.UpdateTrigger.", + "type": "object", + "properties": { + "trigger": { + "description": "Required. The trigger to update.", + "$ref": "#/components/schemas/TriggerUpdate" + }, + "id": { + "description": "Required. The ID of the trigger to update.", + "type": "string" + }, + "update_mask": { + "description": "Optional. The update mask applies to the resource.", + "type": "string", + "format": "google-fieldmask", + "pattern": "^(\\s*[^,\\s.]+(\\s*[,.]\\s*[^,\\s.]+)*)?$" + } + }, + "required": [ + "id", + "trigger" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "UpdateWebhookRequest": { + "description": "Request message for WebhookService.UpdateWebhook.", + "type": "object", + "properties": { + "webhook": { + "description": "Required. The webhook to update.", + "$ref": "#/components/schemas/WebhookUpdate" + }, + "id": { + "description": "Required. The ID of the webhook to update.", + "type": "string" + }, + "update_mask": { + "description": "Optional. The list of fields to update.", + "type": "string", + "format": "google-fieldmask", + "pattern": "^(\\s*[^,\\s.]+(\\s*[,.]\\s*[^,\\s.]+)*)?$" + } + }, + "required": [ + "id", + "webhook" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "UploadEnvironmentFileRequest": { + "description": "Request for `UploadEnvironmentFile`.", + "type": "object", + "properties": { + "overwrite": { + "description": "Optional. Whether to overwrite the destination file if it already exists.", + "type": "boolean" + }, + "extract": { + "description": "Optional. If true, treats the uploaded file as a tar/tar.gz archive and\nunpacks it into `path`.", + "type": "boolean" + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "UploadEnvironmentFileResponse": { + "description": "Response for `UploadEnvironmentFile`.", + "type": "object", + "properties": { + "files": { + "description": "List of files created or extracted in the environment.", + "type": "array", + "items": { + "$ref": "#/components/schemas/EnvironmentFile" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "UrlCitation": { + "description": "A URL citation annotation.", + "type": "object", + "properties": { + "type": { + "const": "url_citation" + }, + "url": { + "description": "The URL.", + "type": "string" + }, + "title": { + "description": "The title of the URL.", + "type": "string" + }, + "start_index": { + "description": "Start of segment of the response that is attributed to this source.\n\nIndex indicates the start of the segment, measured in bytes.", + "type": "integer", + "format": "int32" + }, + "end_index": { + "description": "End of the attributed segment, exclusive.", + "type": "integer", + "format": "int32" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "URLCitation" + }, + "UrlContext": { + "description": "A tool that can be used by the model to fetch URL context.", + "type": "object", + "properties": { + "type": { + "const": "url_context" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "URLContext", + "x-codeSamples": [ + { + "label": "url_context", + "lang": "sh", + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"model\": \"gemini-flash-latest\",\n \"tools\": [{\n \"type\": \"url_context\"\n }],\n \"input\": \"Summarize https://www.example.com\"\n }'\n" + }, + { + "label": "url_context", + "lang": "python", + "source": "from google import genai\n\nclient = genai.Client()\nresponse = client.interactions.create(\n model=\"gemini-flash-latest\",\n tools=[{\"type\": \"url_context\"}],\n input=\"Summarize https://www.example.com\",\n)\nprint(response.output_text)\n" + }, + { + "label": "url_context", + "lang": "javascript", + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n model: 'gemini-flash-latest',\n tools: [{type: 'url_context'}],\n input: 'Summarize https://www.example.com',\n});\nconsole.log(interaction.output_text);\n" + }, + { + "label": "url_context", + "lang": "java", + "source": "import com.google.genai.Client;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.CreateModelInteraction;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\nimport com.google.genai.gaos.models.interactions.URLContext;\n\nClient client = new Client();\nCreateModelInteraction params =\n CreateModelInteraction.builder()\n .model(\"gemini-flash-latest\")\n .tools(List.of(new URLContext()))\n .input(InteractionsInput.of(\"Summarize https://www.example.com\"))\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.outputText().orElse(\"\"));\n" + }, + { + "label": "url_context", + "lang": "go", + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n tool := gaos_interactions.NewTool(gaos_interactions.URLContext{})\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateModelInteraction{\n Model: gaos_interactions.Model(\"gemini-flash-latest\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"Summarize https://www.example.com\")),\n Tools: []gaos_interactions.Tool{tool},\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil && res.Interaction.OutputText != nil {\n fmt.Println(*res.Interaction.OutputText)\n }\n}\n" + } + ] + }, + "UrlContextCallArguments": { + "description": "The arguments to pass to the URL context.", + "type": "object", + "properties": { + "urls": { + "description": "The URLs to fetch.", + "type": "array", + "items": { + "type": "string" + } + } + }, + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "URLContextCallArguments" + }, + "UrlContextCallDelta": { + "type": "object", + "properties": { + "arguments": { + "$ref": "#/components/schemas/UrlContextCallArguments" + }, + "type": { + "const": "url_context_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "arguments", + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "URLContextCallDelta" + }, + "UrlContextCallStep": { + "description": "URL context call step.", + "type": "object", + "properties": { + "arguments": { + "description": "Required. The arguments to pass to the URL context.", + "$ref": "#/components/schemas/UrlContextCallArguments" + }, + "id": { + "description": "Required. A unique ID for this specific tool call.", + "type": "string" + }, + "type": { + "const": "url_context_call" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "arguments", + "id", + "type" + ], + "examples": [ + { + "url_context_call": { + "summary": "UrlContextCallStep", + "value": { + "type": "url_context_call", + "arguments": { + "urls": [ + "https://www.example.com" + ] + }, + "id": "url_call_10219" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "URLContextCallStep" + }, + "UrlContextResult": { + "description": "The result of the URL context.", + "type": "object", + "properties": { + "url": { + "description": "The URL that was fetched.", + "type": "string" + }, + "status": { + "description": "The status of the URL retrieval.", + "type": "string", + "enum": [ + "success", + "error", + "paywall", + "unsafe" + ], + "x-google-enum-descriptions": [ + "Url retrieval is successful.", + "Url retrieval is failed due to error.", + "Url retrieval is failed because the content is behind paywall.", + "Url retrieval is failed because the content is unsafe." + ] + } + }, + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "URLContextResult" + }, + "UrlContextResultDelta": { + "type": "object", + "properties": { + "result": { + "type": "array", + "items": { + "$ref": "#/components/schemas/UrlContextResult" + } + }, + "type": { + "const": "url_context_result" + }, + "is_error": { + "type": "boolean" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "result", + "type" + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "URLContextResultDelta" + }, + "UrlContextResultStep": { + "description": "URL context result step.", + "type": "object", + "properties": { + "result": { + "description": "Required. The results of the URL context.", + "type": "array", + "items": { + "$ref": "#/components/schemas/UrlContextResult" + } + }, + "call_id": { + "description": "Required. ID to match the ID from the function call block.", + "type": "string" + }, + "type": { + "const": "url_context_result" + }, + "is_error": { + "description": "Whether the URL context resulted in an error.", + "type": "boolean" + }, + "signature": { + "description": "A signature hash for backend validation.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + } + }, + "required": [ + "call_id", + "result", + "type" + ], + "examples": [ + { + "url_context_result": { + "summary": "UrlContextResultStep", + "value": { + "type": "url_context_result", + "call_id": "url_call_10219", + "result": [ + { + "title": "Example Domain", + "url": "https://www.example.com", + "snippet": "This domain is for use in illustrative examples in documents." + } + ] + } + } + } + ], + "x-speakeasy-model-namespace": "interactions", + "x-speakeasy-name-override": "URLContextResultStep" + }, + "Usage": { + "description": "Statistics on the interaction request's token usage.", + "type": "object", + "properties": { + "total_input_tokens": { + "description": "Number of tokens in the prompt (context).", + "type": "integer", + "format": "int32" + }, + "input_tokens_by_modality": { + "description": "A breakdown of input token usage by modality.", + "type": "array", + "items": { + "$ref": "#/components/schemas/ModalityTokens" + } + }, + "total_cached_tokens": { + "description": "Number of tokens in the cached part of the prompt (the cached content).", + "type": "integer", + "format": "int32" + }, + "cached_tokens_by_modality": { + "description": "A breakdown of cached token usage by modality.", + "type": "array", + "items": { + "$ref": "#/components/schemas/ModalityTokens" + } + }, + "total_output_tokens": { + "description": "Total number of tokens across all the generated responses.", + "type": "integer", + "format": "int32" + }, + "output_tokens_by_modality": { + "description": "A breakdown of output token usage by modality.", + "type": "array", + "items": { + "$ref": "#/components/schemas/ModalityTokens" + } + }, + "total_tool_use_tokens": { + "description": "Number of tokens present in tool-use prompt(s).", + "type": "integer", + "format": "int32" + }, + "tool_use_tokens_by_modality": { + "description": "A breakdown of tool-use token usage by modality.", + "type": "array", + "items": { + "$ref": "#/components/schemas/ModalityTokens" + } + }, + "total_thought_tokens": { + "description": "Number of tokens of thoughts for thinking models.", + "type": "integer", + "format": "int32" + }, + "total_tokens": { + "description": "Total token count for the interaction request (prompt + responses + other\ninternal tokens).", + "type": "integer", + "format": "int32" + }, + "grounding_tool_count": { + "description": "Grounding tool count.", + "type": "array", + "items": { + "$ref": "#/components/schemas/GroundingToolCount" + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "UserInputStep": { + "description": "Input provided by the user.", + "type": "object", + "properties": { + "type": { + "const": "user_input" + }, + "content": { + "title": "ContentList", + "type": "array", + "items": { + "$ref": "#/components/schemas/Content" + } + } + }, + "required": [ + "type" + ], + "examples": [ + { + "user_input": { + "summary": "UserInputStep", + "value": { + "type": "user_input", + "content": [ + { + "type": "text", + "text": "What is the capital of France?" + } + ] + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "VerbatimTranscriptionMode": { + "description": "Configuration for verbatim transcription mode.", + "type": "object", + "properties": { + "type": { + "const": "verbatim" + }, + "timestamp_granularities": { + "description": "Optional. The granularity of timestamps to include in the transcription output.\nSupported values: \"word\". If empty, no timestamps are generated.", + "type": "array", + "items": { + "type": "string" + } + }, + "diarization_mode": { + "description": "Optional. Configures speaker diarization. Supported values: \"speaker\".", + "type": "string" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "VideoConfig": { + "description": "Configuration options for video generation.", + "type": "object", + "properties": { + "task": { + "description": "Optional task mode for video generation. If not specified, the model\nautomatically determines the appropriate mode based on the provided text\nprompt and input media.", + "type": "string", + "enum": [ + "text_to_video", + "image_to_video", + "reference_to_video", + "edit", + "extend" + ], + "x-google-enum-descriptions": [ + "Generates video solely from a text prompt.", + "Generates video from one or two source images. The first image defines\nthe starting frame, and the optional second image defines the ending\nframe.", + "Generates video using reference media (such as images, audio, or video).", + "Modifies an existing input video.", + "Extends an existing input video." + ] + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "VideoContent": { + "description": "A video content block.", + "type": "object", + "properties": { + "type": { + "const": "video" + }, + "data": { + "description": "The video content.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "description": "The URI of the video.", + "type": "string" + }, + "processing": { + "description": "How the model processes this video for understanding.", + "oneOf": [ + { + "$ref": "#/components/schemas/MediaProcessing" + }, + { + "type": "string", + "enum": [ + "static", + "agentic" + ], + "x-google-enum-descriptions": [ + "Fixed-rate frame extraction. All frames placed in context.", + "Model-driven dynamic navigation." + ] + } + ] + }, + "resolution": { + "description": "The resolution of the media.", + "$ref": "#/components/schemas/MediaResolution" + }, + "name": { + "description": "A user-defined name for this content block. Can be referenced by the model\nin the final response.", + "type": "string" + }, + "mime_type": { + "description": "The mime type of the video.", + "type": "string", + "enum": [ + "video/mp4", + "video/mpeg", + "video/mpg", + "video/mov", + "video/avi", + "video/x-flv", + "video/webm", + "video/wmv", + "video/3gpp" + ], + "x-google-enum-descriptions": [ + "MP4 video format", + "MPEG video format", + "MPG video format", + "MOV video format", + "AVI video format", + "FLV video format", + "WebM video format", + "WMV video format", + "3GPP video format" + ] + } + }, + "required": [ + "type" + ], + "examples": [ + { + "video": { + "summary": "Video", + "value": { + "type": "video", + "uri": "https://www.youtube.com/watch?v=9hE5-98ZeCg" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "VideoDelta": { + "type": "object", + "properties": { + "type": { + "const": "video" + }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, + "mime_type": { + "type": "string", + "enum": [ + "video/mp4", + "video/mpeg", + "video/mpg", + "video/mov", + "video/avi", + "video/x-flv", + "video/webm", + "video/wmv", + "video/3gpp", + "video/jpeg2000" + ], + "x-google-enum-descriptions": [ + "MP4 video format", + "MPEG video format", + "MPG video format", + "MOV video format", + "AVI video format", + "FLV video format", + "WebM video format", + "WMV video format", + "3GPP video format", + "JPEG 2000 video format" + ] + }, + "resolution": { + "description": "The resolution of the media.", + "$ref": "#/components/schemas/MediaResolution" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "VideoResponseFormat": { + "description": "Configuration for video output format.", + "type": "object", + "properties": { + "type": { + "const": "video" + }, + "delivery": { + "description": "The delivery mode for the video output.", + "type": "string", + "enum": [ + "inline", + "uri" + ], + "x-google-enum-descriptions": [ + "Video data is returned inline in the response.", + "Video data is returned as a URI." + ] + }, + "aspect_ratio": { + "description": "The aspect ratio for the video output.", + "type": "string", + "enum": [ + "16:9", + "9:16" + ], + "x-google-enum-descriptions": [ + "16:9 aspect ratio.", + "9:16 aspect ratio." + ] + }, + "duration": { + "description": "The duration for the video output.", + "type": "string", + "format": "google-duration" + }, + "resolution": { + "description": "The video output resolution. Defaults to 720p.", + "type": "string", + "enum": [ + "360p", + "720p", + "1080p", + "4k" + ], + "x-google-enum-descriptions": [ + "360p resolution.", + "720p resolution.", + "1080p resolution.", + "4K resolution." + ] + } + }, + "required": [ + "type" + ], + "examples": [ + { + "video_response_format": { + "summary": "Video Output", + "value": { + "type": "video", + "aspect_ratio": "16:9", + "delivery": "inline" + } + } + } + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Voice": { + "description": "A voice resource representing either a custom voice (created via\n`CreateVoice`) or a prebuilt system voice (returned by `ListVoices`).", + "type": "object", + "properties": { + "type": { + "description": "Required. The type of the voice. In `CreateVoice`, must be `\"replicated\"`\nor `\"prompted\"`. In `ListVoices` responses, may also be `\"prebuilt\"`.", + "$ref": "#/components/schemas/VoiceType" + }, + "id": { + "description": "Output only. The unique identifier of the voice.\n\n* For Google-managed custom voices (`store = true`), this is a generated ID\n with prefix `voice_` (for example, `voice_abc123def456`). Pass\n `voices/{id}` as the `name` in `GetVoice` and `DeleteVoice`, and pass\n `{id}` directly to `SpeechConfig.voice_config.voice` (or\n `SpeechConfig.voice`) during speech synthesis.\n* For prebuilt catalog voices (`\"prebuilt\"` returned by `ListVoices`), this\n is the speaker name (for example, `Puck` or `Charon`).\n* Unset when `CreateVoice` is called with `store = false`.", + "type": "string", + "readOnly": true + }, + "model": { + "description": "Optional. The model used to design or replicate the voice.\nIf omitted in `CreateVoice`, defaults to the latest supported voice design\nmodel. Returned in `CreateVoice`, `GetVoice`, and `ListVoices` responses\nfor custom voices (`\"prompted\"` and `\"replicated\"`); unset for `\"prebuilt\"`\nvoices. Created voices can be synthesized across any supported TTS\nsynthesis model.", + "type": "string" + }, + "key": { + "description": "Output only. The client-managed voice replication key (with prefix\n`voicekey_`). Returned only by `CreateVoice` when `type` is `\"replicated\"`\nand `store` is `false`. Pass this key to `SpeechConfig.voice_config.voice`\n(or `SpeechConfig.voice`) during speech synthesis.", + "type": "string", + "readOnly": true + }, + "replicated": { + "description": "Input only. Parameters for replicated voice generation.\nRequired on input when `type` is `\"replicated\"`. Not returned in\nresponses.", + "$ref": "#/components/schemas/ReplicatedVoice", + "writeOnly": true + }, + "expire_time": { + "description": "Output only. The timestamp at which a custom stored voice (`store = true`)\nor replicated voice key (`store = false`) expires. For custom stored voices\n(`store = true`), this expiration time is extended when the voice is used\nfor speech synthesis or as a `base_voice` in `CreateVoice`. Unset for\nprebuilt catalog voices (`\"prebuilt\"`), which do not expire.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "display_name": { + "description": "Optional. User-provided display name for a stored voice (`store = true`),\nor the catalog name for a prebuilt voice.", + "type": "string" + }, + "prompted": { + "description": "Parameters for prompted voice generation.\nRequired in `CreateVoice` when `type` is `\"prompted\"`. Returned in\n`CreateVoice`, `GetVoice`, and `ListVoices` responses for prompted voices.", + "$ref": "#/components/schemas/PromptedVoice" + }, + "language_code": { + "description": "Optional. Primary BCP-47 language tag (e.g. \"en-US\", \"fr-FR\").", + "type": "string" + }, + "region_code": { + "description": "Optional. ISO 3166-1 alpha-2 or UN M.49 geographic region code (e.g. \"US\",\n\"GB\", \"001\").", + "type": "string" + }, + "accent": { + "description": "Optional. Regional accent descriptor (e.g. \"American\", \"British\").", + "type": "string" + }, + "persona": { + "description": "Optional. Intended persona or character archetype (e.g. \"Warm, Friendly\",\n\"Narrator\").", + "type": "string" + }, + "context": { + "description": "Optional. Optimal usage context or domain (e.g. \"Conversational\",\n\"Audiobook\", \"News\").", + "type": "string" + }, + "gender": { + "description": "Optional. Perceived voice gender presentation (e.g. \"female\", \"male\",\n\"neutral\").", + "type": "string" + }, + "pitch": { + "description": "Optional. Voice pitch classification.", + "$ref": "#/components/schemas/Pitch" + }, + "description": { + "description": "Optional. Descriptive summary of vocal timbre, personality, and tone.", + "type": "string" + }, + "usage": { + "description": "Output only. Token usage statistics for the voice creation request.\nPopulated only in the response of `CreateVoice` when `type` is\n`\"prompted\"`; unset for `\"replicated\"` and in `GetVoice` and `ListVoices`\nresponses.", + "$ref": "#/components/schemas/Usage", + "readOnly": true + }, + "sample_audio": { + "description": "Output only. Sample audio (synthesizer prompt audio) generated for a\nprompted voice.\nPopulated only in `CreateVoice` and `GetVoice` responses when `type` is\n`\"prompted\"`; unset in `ListVoices` responses and for `\"replicated\"` or\n`\"prebuilt\"` voices.", + "$ref": "#/components/schemas/AudioData", + "readOnly": true + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "VoiceType": { + "type": "string", + "enum": [ + "replicated", + "prompted", + "prebuilt" + ], + "x-google-enum-descriptions": [ + "A custom voice replicated from a reference audio sample and speaker consent\nrecording.", + "A custom voice generated from a natural-language text prompt (optionally\nediting or remixing a `base_voice`).", + "A built-in system voice from Google's voice catalog (e.g., `Puck`,\n`Charon`). Returned in responses; cannot be specified as the type in\n`CreateVoice`." + ], + "x-speakeasy-model-namespace": "interactions" + }, + "Webhook": { + "description": "A Webhook resource.", + "type": "object", + "properties": { + "uri": { + "description": "Required. The URI to which webhook events will be sent.", + "type": "string" + }, + "subscribed_events": { + "description": "Required. The events that the webhook is subscribed to.\nAvailable events:\n- batch.succeeded\n- batch.expired\n- batch.failed\n- interaction.requires_action\n- interaction.completed\n- interaction.failed\n- video.generated", + "type": "array", + "items": { + "type": "string" + } + }, + "name": { + "description": "Optional. The user-provided name of the webhook.", + "type": "string" + }, + "create_time": { + "description": "Output only. The timestamp when the webhook was created.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "update_time": { + "description": "Output only. The timestamp when the webhook was last updated.", + "type": "string", + "format": "date-time", + "readOnly": true + }, + "signing_secrets": { + "description": "Output only. The signing secrets associated with this webhook.", + "type": "array", + "items": { + "$ref": "#/components/schemas/SigningSecret" + }, + "readOnly": true + }, + "state": { + "description": "Output only. The state of the webhook.", + "type": "string", + "enum": [ + "enabled", + "disabled", + "disabled_due_to_failed_deliveries" + ], + "readOnly": true, + "x-google-enum-descriptions": [ + "The webhook is enabled.", + "The webhook is disabled by the user.", + "The webhook is disabled due to failed deliveries." + ] + }, + "new_signing_secret": { + "description": "Output only. The new signing secret for the webhook. Only populated on create.", + "type": "string", + "readOnly": true + }, + "id": { + "description": "Output only. The ID of the webhook.", + "type": "string", + "readOnly": true + } + }, + "required": [ + "subscribed_events", + "uri" + ], + "x-speakeasy-model-namespace": "interactions" + }, + "WebhookConfig": { + "description": "Message for configuring webhook events for a request.", + "type": "object", + "properties": { + "uris": { + "description": "Optional. If set, these webhook URIs will be used for webhook events instead of the\nregistered webhooks.", + "type": "array", + "items": { + "type": "string" + } + }, + "user_metadata": { + "description": "Optional. The user metadata that will be returned on each event emission to the\nwebhooks.", + "type": "object", + "additionalProperties": { + "description": "Properties of the object." + } + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "WebhookUpdate": { + "type": "object", + "properties": { + "name": { + "description": "Optional. The user-provided name of the webhook.", + "type": "string" + }, + "uri": { + "description": "Optional. The URI to which webhook events will be sent.", + "type": "string" + }, + "subscribed_events": { + "description": "Optional. The events that the webhook is subscribed to.\nAvailable events:\n- batch.succeeded\n- batch.expired\n- batch.failed\n- interaction.requires_action\n- interaction.completed\n- interaction.failed\n- video.generated", + "type": "array", + "items": { + "type": "string" + } + }, + "state": { + "description": "Optional. The state of the webhook.", + "type": "string", + "enum": [ + "enabled", + "disabled", + "disabled_due_to_failed_deliveries" + ], + "x-google-enum-descriptions": [ + "The webhook is enabled.", + "The webhook is disabled by the user.", + "The webhook is disabled due to failed deliveries." + ] + } + }, + "x-speakeasy-model-namespace": "interactions" + }, + "WordInfo": { + "description": "Word-level ASR annotation for transcription output.\nCarries the word text, optional timing, and optional speaker attribution.", + "type": "object", + "properties": { + "type": { + "const": "word_info" + }, + "text": { + "description": "The transcribed word.", + "type": "string" + }, + "start_offset": { + "description": "Start offset in time of the word relative to the start of the audio.\nPresent when timestamp_granularities contains \"word\".", + "type": "string", + "format": "google-duration" + }, + "end_offset": { + "description": "End offset in time of the word relative to the start of the audio.\nPresent when timestamp_granularities contains \"word\".", + "type": "string", + "format": "google-duration" + }, + "speaker": { + "description": "Optional. Speaker label for this word (e.g. \"spk_1\", \"spk_2\").\nPresent when diarization_mode is set in TranscriptionConfig.", + "type": "string" + }, + "start_index": { + "description": "Start of segment of the response that is attributed to this source.\n\nIndex indicates the start of the segment, measured in bytes.", + "type": "integer", + "format": "int32" + }, + "end_index": { + "description": "End of the attributed segment, exclusive.", + "type": "integer", + "format": "int32" + } + }, + "required": [ + "type" + ], + "x-speakeasy-model-namespace": "interactions" + } + }, + "securitySchemes": { + "googleGenAIAuth": { + "type": "http", + "scheme": "custom", + "x-speakeasy-custom-security-scheme": { + "schema": { + "type": "object", + "properties": { + "accessToken": { + "description": "OAuth access token sent as a bearer Authorization header.", + "type": "string" + }, + "apiKey": { + "description": "Gemini API key sent as x-goog-api-key.", + "type": "string" + }, + "defaultHeaders": { + "description": "Additional default headers to apply before request-specific headers and auth.", + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "required": [] + } + } + } + } + }, + "security": [ + { + "googleGenAIAuth": [] + } + ], + "x-speakeasy-globals": { + "parameters": [ + { + "$ref": "#/components/parameters/api_version" + }, + { + "$ref": "#/components/parameters/google_genai_user_project" + } + ] + }, + "x-speakeasy-retries": { + "backoff": { + "exponent": 2, + "initialInterval": 500, + "maxElapsedTime": 30000, + "maxInterval": 8000 + }, + "retryConnectionErrors": true, + "statusCodes": [ + 408, + 409, + 429, + "5XX" + ], + "strategy": "backoff" + } +} diff --git a/tests/unit/interactions/test_openapi_compliance.py b/tests/unit/interactions/test_openapi_compliance.py index 247d02298aa..b6cc92e0adc 100644 --- a/tests/unit/interactions/test_openapi_compliance.py +++ b/tests/unit/interactions/test_openapi_compliance.py @@ -1,41 +1,56 @@ """ OpenAPI compliance tests for Google Interactions API. -Validates that our SDK requests/responses match the OpenAPI spec at: -https://ai.google.dev/static/api/interactions.openapi.json +Validates that our SDK requests/responses match the OpenAPI spec snapshot in +interactions.openapi.json, taken from https://ai.google.dev/static/api/interactions.openapi.json +(litellm/types/interactions/README.md says how to refresh it). The scheduled +interactions-spec-canary workflow runs the same tests against the live spec by passing +--interactions-spec-url, so a change Google makes shows up there, not in every PR. Run with: pytest tests/unit/interactions/test_openapi_compliance.py -v """ import json -import os import re -from typing import Any, Dict +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any, Dict, Final, get_args from unittest.mock import MagicMock, patch import httpx import pytest from openapi_core import OpenAPI +from pydantic import BaseModel, TypeAdapter OPENAPI_SPEC_URL = "https://ai.google.dev/static/api/interactions.openapi.json" +OPENAPI_SPEC_SNAPSHOT: Final = Path(__file__).with_name("interactions.openapi.json") +SPEC_URL_OPTION: Final = TypeAdapter(str | None) -def _load_openapi_spec_dict() -> Dict[str, Any]: - """ - Load the OpenAPI spec JSON. +@dataclass(frozen=True, slots=True) +class LoadedSpec: + source: str + spec: Dict[str, Any] - In CI or offline environments, network access may not be available. - In that case, gracefully skip these tests instead of erroring. - """ - try: - response = httpx.get(OPENAPI_SPEC_URL, timeout=5.0) - response.raise_for_status() - return response.json() - except Exception as e: # pragma: no cover - defensive, env-dependent - pytest.skip( - f"Skipping Google Interactions OpenAPI compliance tests - " - f"unable to load spec from {OPENAPI_SPEC_URL}: {e}" - ) + +def _status_enum_of(model: type[BaseModel]) -> type[Enum]: + annotation: Final = model.model_fields["status"].annotation + candidates: Final = (annotation, *get_args(annotation)) + return next(candidate for candidate in candidates if isinstance(candidate, type) and issubclass(candidate, Enum)) + + +def _requested_spec_url(request: pytest.FixtureRequest) -> str | None: + return SPEC_URL_OPTION.validate_python(request.config.getoption("interactions_spec_url")) + + +def _load_openapi_spec(live_url: str | None) -> LoadedSpec: + if live_url is None: + return LoadedSpec(source=str(OPENAPI_SPEC_SNAPSHOT), spec=json.loads(OPENAPI_SPEC_SNAPSHOT.read_text())) + with httpx.Client(transport=httpx.HTTPTransport(retries=3), timeout=30.0) as client: + response: Final = client.get(live_url) + response.raise_for_status() + return LoadedSpec(source=str(response.url), spec=response.json()) def _model_create_request_schema(spec_dict: Dict[str, Any]) -> Dict[str, Any]: @@ -65,9 +80,13 @@ def _declared_type_value(variant_schema: Dict[str, Any]) -> Any: @pytest.fixture(scope="module") -def spec_dict() -> Dict[str, Any]: - """Load raw spec dict for manual validation.""" - return _load_openapi_spec_dict() +def loaded_spec(request: pytest.FixtureRequest) -> LoadedSpec: + return _load_openapi_spec(_requested_spec_url(request)) + + +@pytest.fixture(scope="module") +def spec_dict(loaded_spec: LoadedSpec) -> Dict[str, Any]: + return loaded_spec.spec @pytest.fixture(scope="module") @@ -76,6 +95,18 @@ def openapi_spec(spec_dict: Dict[str, Any]) -> OpenAPI: return OpenAPI.from_dict(spec_dict) +def test_spec_comes_from_the_requested_source(loaded_spec: LoadedSpec, request: pytest.FixtureRequest) -> None: + assert loaded_spec.source == (_requested_spec_url(request) or str(OPENAPI_SPEC_SNAPSHOT)) + + +def test_interaction_status_export_carries_the_interaction_models_statuses() -> None: + from litellm.types.interactions import InteractionStatus, generated + + exported: Final = frozenset(member.value for member in InteractionStatus) + modelled: Final = frozenset(member.value for member in _status_enum_of(generated.Interaction)) + assert exported == modelled, sorted(exported ^ modelled) + + class TestRequestCompliance: """Tests that our request bodies match the OpenAPI spec.""" @@ -249,27 +280,24 @@ class TestResponseCompliance: assert field in schema["properties"], f"Output field '{field}' not in spec" print(f"✓ Output field '{field}' exists in spec") - def test_status_enum_values(self, spec_dict): - """Verify status enum values match spec.""" - # `status` is an output-only field; validate against the response schema. - schema = spec_dict["components"]["schemas"]["Interaction"] - status_prop = schema["properties"]["status"] - # Google Interactions API uses lowercase status values (updated Feb 2026). - # Keep this an exact match: this test intentionally breaks CI when - # Google changes the live spec — that breakage is how we get notified - # to review the change. - expected_statuses = [ - "in_progress", - "requires_action", - "completed", - "failed", - "cancelled", - "incomplete", - "budget_exceeded", - "queued", - ] - assert status_prop["enum"] == expected_statuses - print(f"✓ Status enum values: {expected_statuses}") + @pytest.mark.parametrize( + "schema_name", + [ + "Interaction", + "InteractionStatusUpdate", + "CreateModelInteractionParams", + "CreateAgentInteractionParams", + "UrlContextResult", + ], + ) + def test_status_enum_values(self, spec_dict, schema_name): + from litellm.types.interactions import generated + + schema = spec_dict["components"]["schemas"][schema_name] + spec_statuses: Final = frozenset(schema["properties"]["status"]["enum"]) + parsed_statuses: Final = frozenset(member.value for member in _status_enum_of(getattr(generated, schema_name))) + assert spec_statuses + assert spec_statuses <= parsed_statuses, sorted(spec_statuses - parsed_statuses) def test_usage_schema(self, spec_dict): """Verify Usage schema fields.""" diff --git a/tests/unit/litellm_core_utils/attribute_forwarder.py b/tests/unit/litellm_core_utils/attribute_forwarder.py new file mode 100644 index 00000000000..3056057fc97 --- /dev/null +++ b/tests/unit/litellm_core_utils/attribute_forwarder.py @@ -0,0 +1,6 @@ +class AttributeForwarder: + def __init__(self, target: object) -> None: + self._target = target + + def __getattr__(self, name: str) -> object: + return getattr(self._target, name) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index c99147d2e63..21b527250c4 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -65,8 +65,8 @@ def test_missing_cache_read_rate_resolves_to_input_rate(prompt_tokens, read_rate prompt_cost, _ = generic_cost_per_token( "policy-fixture", usage, "openai", service_tier=service_tier, model_info=info ) - assert billed[4] == pytest.approx(read_rate if read_rate is not None else billed[0]) - assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4]) + assert billed.cache_read_rate == pytest.approx(read_rate if read_rate is not None else billed.input_rate) + assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed.input_rate + 100 * billed.cache_read_rate) def test_generic_cost_per_token_bills_cache_reads_at_input_rate_when_no_cache_read_rate() -> None: @@ -248,7 +248,7 @@ def test_missing_cache_read_uses_off_peak_input_rate(): } when = datetime(2026, 9, 7, 12, tzinfo=timezone.utc) billed = _get_token_base_cost(info, Usage(prompt_tokens=100), current_time=when) - assert billed[0] == billed[4] == 5e-6 + assert billed.input_rate == billed.cache_read_rate == 5e-6 def test_reasoning_tokens_no_price_set(_local_model_cost_map): @@ -324,8 +324,6 @@ def test_reasoning_tokens_gemini(_local_model_cost_map): ) - - def test_image_tokens_with_custom_pricing(): """Test that image_tokens in completion are properly costed with output_cost_per_image_token.""" from unittest.mock import patch @@ -578,7 +576,7 @@ def test_get_token_base_cost_picks_highest_crossed_tier(): } usage = Usage(prompt_tokens=150_000, completion_tokens=10, total_tokens=150_010) - prompt_base_cost = _get_token_base_cost(model_info, usage)[0] + prompt_base_cost = _get_token_base_cost(model_info, usage).input_rate assert prompt_base_cost == 9e-6 @@ -779,12 +777,12 @@ def test_get_token_base_cost_weekend_only_off_peak_rate(): saturday_peak_hours = _get_token_base_cost( model_info, usage, current_time=datetime(2026, 8, 29, 2, 0, tzinfo=timezone.utc) ) - assert saturday_peak_hours[:2] == (5e-7, 1e-6) + assert (saturday_peak_hours.input_rate, saturday_peak_hours.output_rate) == (5e-7, 1e-6) monday_same_hours = _get_token_base_cost( model_info, usage, current_time=datetime(2026, 8, 24, 2, 0, tzinfo=timezone.utc) ) - assert monday_same_hours[:2] == (1e-6, 2e-6) + assert (monday_same_hours.input_rate, monday_same_hours.output_rate) == (1e-6, 2e-6) def test_get_token_base_cost_applies_off_peak_pricing(): @@ -810,14 +808,14 @@ def test_get_token_base_cost_applies_off_peak_pricing(): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert off_peak[0] == 5e-7 - assert off_peak[1] == 1e-6 - assert off_peak[4] == 5e-8 + assert off_peak.input_rate == 5e-7 + assert off_peak.output_rate == 1e-6 + assert off_peak.cache_read_rate == 5e-8 peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert peak[0] == 1e-6 - assert peak[1] == 2e-6 - assert peak[4] == 1e-7 + assert peak.input_rate == 1e-6 + assert peak.output_rate == 2e-6 + assert peak.cache_read_rate == 1e-7 def test_get_token_base_cost_non_mapping_off_peak_block_bills_standard_rates(): @@ -843,8 +841,8 @@ def test_get_token_base_cost_non_mapping_off_peak_block_bills_standard_rates(): }, ) result = _get_token_base_cost(model_info, usage, current_time=when) - assert result[0] == 1e-6 - assert result[1] == 2e-6 + assert result.input_rate == 1e-6 + assert result.output_rate == 2e-6 def test_get_token_base_cost_off_peak_falls_back_to_standard_when_unset(): @@ -864,8 +862,8 @@ def test_get_token_base_cost_off_peak_falls_back_to_standard_when_unset(): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) result = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert result[0] == 5e-7 - assert result[1] == 2e-6 + assert result.input_rate == 5e-7 + assert result.output_rate == 2e-6 def test_get_token_base_cost_off_peak_wins_over_threshold(): @@ -891,12 +889,12 @@ def test_get_token_base_cost_off_peak_wins_over_threshold(): usage = Usage(prompt_tokens=250000, completion_tokens=250000, total_tokens=500000) off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert off_peak[0] == 5e-7 - assert off_peak[1] == 1e-6 + assert off_peak.input_rate == 5e-7 + assert off_peak.output_rate == 1e-6 peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert peak[0] == 3e-6 - assert peak[1] == 4e-6 + assert peak.input_rate == 3e-6 + assert peak.output_rate == 4e-6 def test_get_model_info_propagates_off_peak_fields(): @@ -948,10 +946,10 @@ def test_get_token_base_cost_off_peak_wins_over_tiered_pricing(): usage = Usage(prompt_tokens=1_000, completion_tokens=100, total_tokens=1_100) inside = _get_token_base_cost(info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert inside[:2] == (5e-7, 1e-6) + assert (inside.input_rate, inside.output_rate) == (5e-7, 1e-6) outside = _get_token_base_cost(info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert outside[:2] == (3e-6, 6e-6) + assert (outside.input_rate, outside.output_rate) == (3e-6, 6e-6) def _register_off_peak_reasoning_model( @@ -1153,17 +1151,17 @@ def test_get_token_base_cost_off_peak_cache_creation_rate(): inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) inside = _get_token_base_cost(model_info, usage, current_time=inside_window) - assert inside[2] == 5e-7 - assert inside[3] == 2e-6 + assert inside.cache_creation_rate == 5e-7 + assert inside.cache_creation_rates_by_ttl["1hr"] == 2e-6 outside = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert outside[2] == 1.25e-6 + assert outside.cache_creation_rate == 1.25e-6 without_key = cast( ModelInfo, {**model_info, "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}}, ) - assert _get_token_base_cost(without_key, usage, current_time=inside_window)[2] == 1.25e-6 + assert _get_token_base_cost(without_key, usage, current_time=inside_window).cache_creation_rate == 1.25e-6 def test_get_token_type_cost_breakdown_reflects_off_peak_reasoning_and_cache_creation_rates(): @@ -1583,19 +1581,14 @@ def test_generic_cost_per_token_tier_1hr_cache_rate_preserves_zero_and_fallback( ), ) costs: Final = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openrouter") - breakdown: Final = get_token_type_cost_breakdown( - model=model, usage=usage, custom_llm_provider="openrouter" - ) + breakdown: Final = get_token_type_cost_breakdown(model=model, usage=usage, custom_llm_provider="openrouter") ordinary_rate: Final = cache_rates.get("cache_creation_input_token_cost") - expected_write_cost: Final = ( - (20 - one_hour_tokens) * (ordinary_rate if ordinary_rate is not None else 1.0) - + one_hour_tokens * expected_write_rate - ) + expected_write_cost: Final = (20 - one_hour_tokens) * ( + ordinary_rate if ordinary_rate is not None else 1.0 + ) + one_hour_tokens * expected_write_rate assert costs == pytest.approx((40 * 1.0 + 40 * 0.1 + expected_write_cost, 10 * 2.0)) assert breakdown.cache_creation_cost == pytest.approx(expected_write_cost) - assert breakdown.rates is not None - assert breakdown.rates.cache_creation_input_token_cost_above_1hr == expected_write_rate finally: litellm.model_cost.pop(model, None) @@ -2020,8 +2013,8 @@ def test_calculate_cache_writing_cost(): result = calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, cache_creation_cost=cache_creation_cost, + cache_creation_cost_by_ttl={"1hr": cache_creation_cost_above_1hr, "12hr": cache_creation_cost_above_1hr}, ) # Expected calculation: @@ -2041,8 +2034,8 @@ def test_calculate_cache_writing_cost(): result_no_details = calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens_no_details, cache_creation_token_details=cache_creation_token_details_none, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, cache_creation_cost=cache_creation_cost_fallback, + cache_creation_cost_by_ttl={"1hr": cache_creation_cost_above_1hr, "12hr": cache_creation_cost_above_1hr}, ) # Expected calculation when no details: 1000 * 5e-06 = 0.005 @@ -2059,15 +2052,13 @@ def test_calculate_cache_writing_cost(): result_partial = calculate_cache_writing_cost( cache_creation_tokens=500, cache_creation_token_details=cache_creation_token_details_partial, - cache_creation_cost_above_1hr=6e-06, cache_creation_cost=3e-06, + cache_creation_cost_by_ttl={"1hr": 6e-06, "12hr": 6e-06}, ) - # Expected calculation: 0 (for None 5m tokens) + (100 * 6e-06) = 0.0006 - expected_cost_partial = (0.0) + (100 * 6e-06) + expected_cost_partial = (400 * 3e-06) + (100 * 6e-06) assert round(result_partial, 6) == round(expected_cost_partial, 6) - assert round(result_partial, 6) == 0.0006 # Test case 4: Zero costs result_zero = calculate_cache_writing_cost( @@ -2075,13 +2066,161 @@ def test_calculate_cache_writing_cost(): cache_creation_token_details=CacheCreationTokenDetails( ephemeral_5m_input_tokens=50, ephemeral_1h_input_tokens=950 ), - cache_creation_cost_above_1hr=0.0, cache_creation_cost=0.0, + cache_creation_cost_by_ttl={"1hr": 0.0, "12hr": 0.0}, ) assert result_zero == 0.0 +@pytest.mark.parametrize("field", CacheCreationTokenDetails.model_fields) +def test_calculate_cache_writing_cost_bills_every_reported_lifetime(field: str) -> None: + cost = calculate_cache_writing_cost( + cache_creation_tokens=10, + cache_creation_token_details=CacheCreationTokenDetails(**{field: 10}), + cache_creation_cost=2.0, + cache_creation_cost_by_ttl={"1hr": 2.0, "12hr": 2.0}, + ) + assert cost == 20.0 + + +def test_calculate_cache_writing_cost_bills_writes_of_an_unknown_lifetime_at_the_base_rate() -> None: + cost = calculate_cache_writing_cost( + cache_creation_tokens=10, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_30m_input_tokens=2, ephemeral_12h_input_tokens=3 + ), + cache_creation_cost=1.0, + cache_creation_cost_by_ttl={"1hr": 4.0, "12hr": 4.0}, + ) + assert cost == 2 * 1.0 + 3 * 4.0 + 5 * 1.0 + + +@pytest.mark.parametrize( + ("prompt_tokens", "service_tier", "write_rate", "twelve_hour_rate"), + [ + (10_000, "default", 2e-6, 4e-6), + (10_000, "flex", 1e-6, 2e-6), + (10_000, "priority", 3e-6, 6e-6), + (300_000, "default", 4e-6, 8e-6), + (300_000, "flex", 2e-6, 4e-6), + (300_000, "priority", 6e-6, 12e-6), + ], +) +def test_generic_cost_per_token_prices_12h_cache_writes_at_prompt_threshold( + prompt_tokens: int, service_tier: str, write_rate: float, twelve_hour_rate: float +) -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost": 2e-6, + "cache_creation_input_token_cost_flex": 1e-6, + "cache_creation_input_token_cost_priority": 3e-6, + "cache_creation_input_token_cost_above_272k_tokens": 4e-6, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2e-6, + "cache_creation_input_token_cost_above_272k_tokens_priority": 6e-6, + "cache_creation_input_token_cost_above_12hr": 4e-6, + "cache_creation_input_token_cost_above_12hr_flex": 2e-6, + "cache_creation_input_token_cost_above_12hr_priority": 6e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 8e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 4e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 12e-6, + } + usage: Final = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=0, + total_tokens=prompt_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cache_creation_tokens=5_000, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_30m_input_tokens=3_000, ephemeral_12h_input_tokens=2_000 + ), + ), + ) + + result: Final = generic_cost_per_token( + model="test-12h-cache-write", + usage=usage, + custom_llm_provider="openai", + service_tier=None if service_tier == "default" else service_tier, + model_info=model_info, + ) + + assert result[0] == pytest.approx( + (usage.prompt_tokens - usage.prompt_tokens_details.cache_creation_tokens) * model_info["input_cost_per_token"] + + 3_000 * write_rate + + 2_000 * twelve_hour_rate + ) + + +def test_generic_cost_per_token_falls_back_to_default_for_missing_12h_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost": 3e-6, + } + usage: Final = Usage( + prompt_tokens=100, + completion_tokens=0, + total_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper( + cache_creation_tokens=100, + cache_creation_token_details=CacheCreationTokenDetails(ephemeral_12h_input_tokens=100), + ), + ) + + result: Final = generic_cost_per_token( + model="test-12h-cache-write-fallback", + usage=usage, + custom_llm_provider="openai", + model_info=model_info, + ) + + assert result[0] == pytest.approx(100 * model_info["cache_creation_input_token_cost"]) + + +@pytest.mark.parametrize(("prompt_tokens", "twelve_hour_rate"), [(10_000, 7e-6), (300_000, 14e-6)]) +def test_generic_cost_per_token_prices_ultrafast_12h_cache_writes_at_the_ultrafast_rate( + prompt_tokens: int, twelve_hour_rate: float +) -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "input_cost_per_token_ultrafast": 5e-6, + "cache_creation_input_token_cost": 2e-6, + "cache_creation_input_token_cost_ultrafast": 3e-6, + "cache_creation_input_token_cost_above_12hr": 4e-6, + "cache_creation_input_token_cost_above_12hr_ultrafast": 7e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 8e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 14e-6, + } + usage: Final = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=0, + total_tokens=prompt_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cache_creation_tokens=5_000, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_30m_input_tokens=3_000, ephemeral_12h_input_tokens=2_000 + ), + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="test-ultrafast-12h-cache-write", + usage=usage, + custom_llm_provider="openai", + service_tier="ultrafast", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx( + (prompt_tokens - 5_000) * model_info["input_cost_per_token_ultrafast"] + + 3_000 * model_info["cache_creation_input_token_cost_ultrafast"] + + 2_000 * twelve_hour_rate + ) + + def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): """ Regression test: when cache_creation_tokens is 0 but cache_creation_token_details @@ -2119,7 +2258,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): prompt_base_cost=0.0, cache_read_cost=0.0, cache_creation_cost=cache_creation_cost, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, + cache_creation_cost_by_ttl={"1hr": cache_creation_cost_above_1hr, "12hr": cache_creation_cost_above_1hr}, ) # Expected: (100 * 3.75e-06) + (200 * 6e-06) = 0.000375 + 0.0012 = 0.001575 @@ -2330,10 +2469,6 @@ def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(_local_mo assert round(cost, 10) == round(expected_cost, 10) - - - - def test_query_count_is_free_without_a_per_query_price(_local_model_cost_map): usage = Usage( prompt_tokens=0, @@ -2512,8 +2647,6 @@ def test_vertex_global_or_absent_location_no_uplift(vertex_location, _local_mode assert base == located - - def test_vertex_uplift_invalid_multiplier_defaults_to_one(): """A malformed multiplier in the cost map degrades to base pricing, never raises.""" from litellm.litellm_core_utils.llm_cost_calc.utils import ( @@ -2569,10 +2702,10 @@ def test_flat_input_price_keeps_other_token_thresholds(field: str) -> None: model_info=info, usage=Usage(prompt_tokens=300_001, completion_tokens=3, total_tokens=300_004), ) - assert rates[0] == 0.001 - assert rates[1] == (0.005 if field == "output_cost_per_token" else 0.002) - assert rates[2] == (0.005 if field == "cache_creation_input_token_cost" else 0.004) - assert rates[4] == (0.005 if field == "cache_read_input_token_cost" else 0.003) + assert rates.input_rate == 0.001 + assert rates.output_rate == (0.005 if field == "output_cost_per_token" else 0.002) + assert rates.cache_creation_rate == (0.005 if field == "cache_creation_input_token_cost" else 0.004) + assert rates.cache_read_rate == (0.005 if field == "cache_read_input_token_cost" else 0.003) @pytest.mark.parametrize("prompt_tokens", (200_000, 200_001, 512_000, 512_001)) @@ -2599,10 +2732,12 @@ def test_mixed_token_thresholds_preserve_each_crossed_rate(prompt_tokens: int, h model_info=info, usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=3, total_tokens=prompt_tokens + 3), ) - assert rates[0] == (0.005 if prompt_tokens > 200_000 else 0.001) - assert rates[1] == (higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002) - assert rates[2] == (0.007 if prompt_tokens > 200_000 else 0.004) - assert rates[4] == (higher_rate if prompt_tokens > 512_000 else 0.003) + assert rates.input_rate == (0.005 if prompt_tokens > 200_000 else 0.001) + assert rates.output_rate == ( + higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002 + ) + assert rates.cache_creation_rate == (0.007 if prompt_tokens > 200_000 else 0.004) + assert rates.cache_read_rate == (higher_rate if prompt_tokens > 512_000 else 0.003) def test_threshold_keys_exclude_service_tier_variants(): @@ -2626,8 +2761,7 @@ def test_threshold_keys_exclude_service_tier_variants(): }, ) usage = Usage(prompt_tokens=350_000, completion_tokens=1_000, total_tokens=351_000) - prompt_base, *_ = _get_token_base_cost(model_info=model_info, usage=usage) - assert prompt_base == 3e-6 + assert _get_token_base_cost(model_info=model_info, usage=usage).input_rate == 3e-6 @pytest.mark.parametrize( @@ -2897,7 +3031,7 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp cache_read_input_token_cost=6e-7, cache_read_input_audio_token_cost=6e-7, cache_creation_input_token_cost=7.5e-6, - cache_creation_input_token_cost_above_1hr=7.5e-6, + cache_creation_input_token_cost_by_ttl={"1hr": 7.5e-6, "12hr": 7.5e-6}, output_cost_per_reasoning_token=3e-5, ) assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost) @@ -3713,6 +3847,23 @@ def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a 3e-7, id="no-1h-price-uses-the-off-peak-write-price", ), + pytest.param( + { + "tiered_pricing": [ + { + "range": [0, 128000], + "input_cost_per_token": 2e-7, + "output_cost_per_token": 1.25e-6, + "cache_creation_input_token_cost": 2.5e-7, + } + ], + "off_peak_pricing": {"hours_utc": "00:00-23:59", "cache_creation_input_token_cost": 3e-7}, + }, + datetime(2026, 9, 14, 12, tzinfo=timezone.utc), + 3e-7, + 3e-7, + id="tiered-no-1h-price-uses-the-off-peak-write-price", + ), ), ) def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_path( @@ -3724,10 +3875,17 @@ def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_ model_info = {"input_cost_per_token": 2e-7, "output_cost_per_token": 1.25e-6, **cache_rates} usage = Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11) - _, _, creation, creation_1h, _ = _get_token_base_cost(model_info, usage, current_time=current_time) + rates = _get_token_base_cost(model_info, usage, current_time=current_time) - assert creation == pytest.approx(expected_creation) - assert creation_1h == pytest.approx(expected_creation_1h) + assert rates.cache_creation_rate == pytest.approx(expected_creation) + assert calculate_cache_writing_cost( + cache_creation_tokens=3, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_1h_input_tokens=1, ephemeral_12h_input_tokens=2 + ), + cache_creation_cost=rates.cache_creation_rate, + cache_creation_cost_by_ttl=rates.cache_creation_rates_by_ttl, + ) == pytest.approx(3 * expected_creation_1h) def _image_response(num_images: int = 1, usage: ImageUsage | None = None) -> ImageResponse: @@ -3787,8 +3945,6 @@ def test_route_image_generation_cost_openai_honors_deployment_input_cost_per_ima assert cost == pytest.approx(0.07) - - @pytest.mark.parametrize( ("custom_llm_provider", "model"), [ @@ -3965,6 +4121,44 @@ def test_get_batch_cost_rates_has_no_cache_write_rate_without_a_cache_write_batc assert rates.cache_creation is None +@pytest.mark.parametrize(("prompt_tokens", "expected"), [(1_000, 2.5e-7), (272_000, 2.5e-7), (300_000, 5e-7)]) +def test_get_batch_cost_rates_reads_the_12h_batch_cache_write_rate_for_the_crossed_tier( + prompt_tokens: int, expected: float +) -> None: + from litellm.litellm_core_utils.llm_cost_calc.utils import get_batch_cost_rates + + rates: Final = get_batch_cost_rates( + _batch_rates_model_info( + input_cost_per_token_batches=1e-7, + cache_creation_input_token_cost_batches=1.25e-7, + cache_creation_input_token_cost_above_12hr_batches=2.5e-7, + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches=5e-7, + ), + Usage(prompt_tokens=prompt_tokens, completion_tokens=1, total_tokens=prompt_tokens + 1), + "openai", + ) + + assert rates.cache_creation == 1.25e-7 + assert rates.cache_creation_rates_by_ttl == {"12hr": expected} + + +def test_get_batch_cost_rates_has_no_12h_cache_write_rate_without_a_12h_batch_key() -> None: + from litellm.litellm_core_utils.llm_cost_calc.utils import get_batch_cost_rates + + rates: Final = get_batch_cost_rates( + _batch_rates_model_info( + input_cost_per_token_batches=1e-7, + cache_creation_input_token_cost_batches=1.25e-7, + cache_creation_input_token_cost_above_12hr=2.5e-7, + cache_creation_input_token_cost_above_12hr_ultrafast=5e-7, + ), + Usage(prompt_tokens=300_000, completion_tokens=1, total_tokens=300_001), + "openai", + ) + + assert rates.cache_creation_rates_by_ttl == {} + + @pytest.mark.parametrize("model_base", ["gpt-6-sol", "gpt-6-luna"]) def test_azure_gpt_6_foundry_price_sheet(_local_model_cost_map, model_base): """Azure Foundry hosts gpt-6-sol and gpt-6-luna at OpenAI's Global rates, with the @@ -4206,17 +4400,28 @@ def test_token_breakdown_uses_supplied_pricing_without_catalog_lookup(rate: floa "cache_creation_input_token_cost": rate, } usage: Final = Usage( - prompt_tokens=20, completion_tokens=3, total_tokens=23, - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), cache_creation_input_tokens=5, + prompt_tokens=20, + completion_tokens=3, + total_tokens=23, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), + cache_creation_input_tokens=5, ) prompt_cost, output_cost = generic_cost_per_token( - model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info, + model="supplied-pricing-test", + usage=usage, + custom_llm_provider="bedrock", + model_info=model_info, ) breakdown: Final = get_token_type_cost_breakdown( - model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info, + model="supplied-pricing-test", + usage=usage, + custom_llm_provider="bedrock", + model_info=model_info, ) assert breakdown.cache_read_cost == pytest.approx(5 * rate) assert breakdown.cache_creation_cost == pytest.approx(5 * rate) assert breakdown.rates is not None - assert prompt_cost == pytest.approx(10 * breakdown.rates.input_cost_per_token + breakdown.cache_read_cost + breakdown.cache_creation_cost) + assert prompt_cost == pytest.approx( + 10 * breakdown.rates.input_cost_per_token + breakdown.cache_read_cost + breakdown.cache_creation_cost + ) assert output_cost == pytest.approx(3 * breakdown.rates.output_cost_per_token) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py index 142d49bfff1..b6c18e0da86 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py @@ -11,6 +11,7 @@ import respx import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.llms.base import CachedTokensDetails from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage TIER_MODEL: Final = "tier-priced-test-model" @@ -31,6 +32,20 @@ PROMPT_TOKENS: Final = 1000 CACHED_TOKENS: Final = 200 COMPLETION_TOKENS: Final = 500 TIER_API_BASE: Final = "https://tier-pricing.invalid/v1" +AUDIO_TIER_MODEL: Final = "tier-priced-gemini-audio-test-model" +AUDIO_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType( + { + "input_cost_per_token": 4e-06, + "output_cost_per_token": 8e-06, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_audio_token": 3e-06, + "cache_read_input_audio_token_cost": 2e-06, + "input_cost_per_audio_token_flex": 1e-06, + "input_cost_per_audio_token_priority": 2e-06, + "cache_read_input_audio_token_cost_flex": 4e-07, + "cache_read_input_audio_token_cost_priority": 8e-07, + } +) def _cost_at(prices: Mapping[str, float], column_suffix: str) -> float: @@ -78,6 +93,53 @@ def test_completion_cost_bills_the_price_columns_of_the_service_tier( assert cost == pytest.approx(_cost_at(TIER_ROW, column_suffix)) +@pytest.mark.parametrize( + ("service_tier", "audio_rate_key", "cache_audio_rate_key"), + ( + pytest.param("flex", "input_cost_per_audio_token_flex", "cache_read_input_audio_token_cost_flex"), + pytest.param("priority", "input_cost_per_audio_token_priority", "cache_read_input_audio_token_cost_priority"), + ), +) +def test_gemini_audio_and_cached_audio_use_service_tier_rates( + local_model_cost_map: None, service_tier: str, audio_rate_key: str, cache_audio_rate_key: str +) -> None: + litellm.register_model( + { + AUDIO_TIER_MODEL: { + "litellm_provider": "gemini", + "mode": "chat", + **AUDIO_TIER_ROW, + } + } + ) + usage: Final = Usage( + prompt_tokens=40, + completion_tokens=0, + total_tokens=40, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, + audio_tokens=20, + cached_tokens=10, + cached_tokens_details=CachedTokensDetails(audio_tokens=10), + ), + ) + response: Final = ModelResponse(model=AUDIO_TIER_MODEL, usage=usage) + + cost: Final = litellm.completion_cost( + completion_response=response, + model=AUDIO_TIER_MODEL, + custom_llm_provider="gemini", + service_tier=service_tier, + ) + + expected: Final = ( + 20 * AUDIO_TIER_ROW["input_cost_per_token"] + + 10 * AUDIO_TIER_ROW[audio_rate_key] + + 10 * AUDIO_TIER_ROW[cache_audio_rate_key] + ) + assert cost == pytest.approx(expected) + + LONG_CONTEXT_TIER_MODEL: Final = "long-context-tier-priced-test-model" LONG_CONTEXT_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType( { diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py index 921dde8bf7b..bdc14c4934b 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import ( used_pricing_keys, zero_cost_warning, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.utils import CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage PER_SECOND_ENTRY: Final = {"cost_per_second": 0.00042} @@ -135,6 +136,36 @@ def test_audio_tokens_need_the_audio_rates() -> None: assert diagnostic["missing_pricing_keys"] == ("input_cost_per_audio_token", "output_cost_per_audio_token") +def test_cached_audio_uses_its_cache_rate_in_zero_cost_diagnostic() -> None: + usage: Final = Usage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=0, + audio_tokens=20, + cached_tokens=20, + cached_tokens_details=CachedTokensDetails(audio_tokens=20), + ), + ) + pricing_entry: Final = { + "input_cost_per_token": 1e-6, + "input_cost_per_audio_token": 2e-6, + "cache_read_input_audio_token_cost": 0, + } + + assert used_pricing_keys(usage) == ("cache_read_input_audio_token_cost",) + assert ( + diagnose_zero_cost( + usage=usage, + pricing_model="gemini-audio", + pricing_entry=pricing_entry, + calculation_failed=False, + ) + is None + ) + + def test_warning_names_the_request_the_entry_the_missing_keys_and_the_counter() -> None: diagnostic = diagnose_zero_cost( usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=False diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 6459aa98e70..d66c37ba8c2 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -1,5 +1,6 @@ """Tests for litellm_core_utils.core_helpers module.""" +import datetime import json import logging import threading @@ -18,13 +19,13 @@ from litellm.litellm_core_utils.core_helpers import ( drop_params_env_flag, drop_params_flag, get_litellm_metadata_from_kwargs, - get_parent_otel_span_from_kwargs, get_or_create_metadata_bucket, + get_parent_otel_span_from_kwargs, get_provider_response_headers_from_hidden_params, map_finish_reason, normalize_drop_params, - process_response_headers, preserve_upstream_non_openai_attributes, + process_response_headers, reconstruct_model_name, redact_nested_match_and_regex_keys, set_provider_response_headers_in_hidden_params, @@ -685,3 +686,36 @@ def test_spend_metadata_merge_reads_request_metadata_that_is_a_read_only_mapping ) assert merged == {"user_api_key_hash": "hash", "model_group": "claude"} + + +class TestReceiveStampSurvivesASecondMetadataBucket: + """A guardrail hook can plant ``litellm_metadata`` on a chat request after the proxy stamped + ``litellm_received_at`` into ``metadata``. The resolved bucket must still carry the stamp, or + ``x-litellm-overhead-duration-ms`` starts at the router call and hides every pre-call cost.""" + + def test_receive_stamp_is_carried_into_litellm_metadata(self): + received_at = datetime.datetime(2026, 10, 21, 5, 0, tzinfo=datetime.timezone.utc) + kwargs = { + "litellm_params": { + "metadata": {"litellm_received_at": received_at, "user_api_key_user_id": "user-1"}, + "litellm_metadata": {"applied_guardrails": ["perf-pre"]}, + } + } + + resolved = get_litellm_metadata_from_kwargs(kwargs) + + assert resolved is kwargs["litellm_params"]["litellm_metadata"] + assert resolved["litellm_received_at"] is received_at + assert resolved["user_api_key_user_id"] == "user-1" + + def test_litellm_metadata_keeps_its_own_receive_stamp(self): + own = datetime.datetime(2026, 10, 21, 5, 0, tzinfo=datetime.timezone.utc) + other = datetime.datetime(2026, 10, 21, 5, 1, tzinfo=datetime.timezone.utc) + kwargs = { + "litellm_params": { + "metadata": {"litellm_received_at": other}, + "litellm_metadata": {"litellm_received_at": own}, + } + } + + assert get_litellm_metadata_from_kwargs(kwargs)["litellm_received_at"] is own diff --git a/tests/unit/litellm_core_utils/test_dot_notation_indexing.py b/tests/unit/litellm_core_utils/test_dot_notation_indexing.py index ae529a71009..7c670f1dac9 100644 --- a/tests/unit/litellm_core_utils/test_dot_notation_indexing.py +++ b/tests/unit/litellm_core_utils/test_dot_notation_indexing.py @@ -2,6 +2,9 @@ Tests for litellm.litellm_core_utils.dot_notation_indexing module. """ +from types import MappingProxyType +from typing import Final + import pytest from litellm.litellm_core_utils.dot_notation_indexing import ( @@ -125,3 +128,14 @@ class TestDeleteNestedValue: data = {"items": [{"a": 1, "b": 2}, {"a": 3, "b": 4}]} result = delete_nested_value(data, "items[0].b") assert result == {"items": [{"a": 1}, {"a": 3, "b": 4}]} + + def test_delete_returns_a_fresh_dict_and_leaves_a_narrowly_typed_input_untouched(self): + payload: Final[dict[str, list[str]]] = {"tags": ["a", "b"], "secret": ["s"]} + result = delete_nested_value(payload, "secret") + assert result == {"tags": ["a", "b"]} + assert payload == {"tags": ["a", "b"], "secret": ["s"]} + assert result["tags"] is not payload["tags"] + + def test_delete_accepts_a_read_only_mapping(self): + payload: Final = MappingProxyType({"tags": ["a"], "secret": ["s"]}) + assert delete_nested_value(payload, "secret") == {"tags": ["a"]} 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 5ee14ff530d..9e8e9cf2a80 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -160,6 +160,12 @@ def test_api_base_line_keeps_an_at_sign_in_the_path(api_base: str) -> None: # Tuple format: (error_message, expected_result) context_window_test_cases = [ # Positive cases (should return True) + ( + '"{\n "error": {\n "message": "Your input exceeds the context window of this model. ' + 'Please adjust your input and try again.",\n "type": "invalid_request_error",\n' + ' "param": "input",\n "code": "context_length_exceeded"\n }\n}"', + True, + ), ( "An error occurred: The input exceeds the model's maximum context limit of 8192 tokens.", True, diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index d7474936425..e9e3f7b8174 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -1,5 +1,9 @@ +from collections.abc import Iterator +from pathlib import Path from typing import Final +import respx + import asyncio, httpx, importlib, os import pytest @@ -8,6 +12,8 @@ from litellm import CustomLLM from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.get_llm_provider_logic import ( get_llm_provider, + resolve_model_and_provider_for_metadata, + get_provider_affinity_header, inferred_provider, is_registered_custom_provider, ) @@ -706,3 +712,55 @@ class TestClaudeModelPatternMatching: set_fallback_generalizations([]) with pytest.raises(litellm.BadRequestError): litellm.get_llm_provider(model="claude-opus-4-9") + + +@pytest.mark.parametrize("prefix", ["openai", "anthropic", "custom_openai", "openai_like"]) +def test_opencode_affinity_resolution_preserves_existing_transport(prefix: str) -> None: + api_base: Final = "https://opencode.ai/zen/go/v1" + model, provider, _, destination = get_llm_provider(model=f"{prefix}/test-model", api_base=api_base) + assert (model, provider, destination) == ("test-model", prefix, api_base) + assert get_provider_affinity_header(destination) == "x-opencode-session" + + +@pytest.fixture +def chatgpt_login_tripwire(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[None]: + blocked: Final = tmp_path / "blocked" + blocked.write_text("ordinary file") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(blocked / "unused")) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + with respx.mock as transport: + yield + assert len(transport.calls) == 0 + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected"), + [ + ("chatgpt/gpt-5.4", None, ("gpt-5.4", "chatgpt")), + ("gpt-5.4", "chatgpt", ("gpt-5.4", "chatgpt")), + ("github_copilot/gpt-4o", None, ("gpt-4o", "github_copilot")), + ("openai/gpt-5.4", None, ("gpt-5.4", "openai")), + ], +) +def test_resolve_model_and_provider_for_metadata_never_logs_in( + chatgpt_login_tripwire: None, model: str, custom_llm_provider: str | None, expected: tuple[str, str] +) -> None: + assert resolve_model_and_provider_for_metadata(model, custom_llm_provider) == expected + + +def test_metadata_callers_never_log_in_for_a_declared_chatgpt_deployment( + chatgpt_login_tripwire: None, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base + from litellm.proxy.spend_tracking.savings import _resolve_model + from litellm.router import Router + from litellm.types.router import LiteLLM_Params + from litellm.utils import supports_native_streaming + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + assert supports_native_streaming("chatgpt/gpt-5.4", None) is True + assert get_api_base("chatgpt/gpt-5.4", {"model": "chatgpt/gpt-5.4"}) == "https://chatgpt.com/backend-api/codex" + identity = _resolve_model("chatgpt/gpt-5.4", None) + assert identity is not None and identity.provider == "chatgpt" + assert Router._model_group_llm_provider(LiteLLM_Params(model="chatgpt/gpt-5.4")) == ("gpt-5.4", "chatgpt") 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 84e785ac99a..d44169ae28a 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 @@ -399,9 +399,10 @@ async def test_refetch_shrunk_map_fails_integrity_not_swapped_in(): @pytest.mark.asyncio -async def test_refetch_respects_local_env_override(monkeypatch): - """LITELLM_LOCAL_MODEL_COST_MAP=True short-circuits to the bundled backup, zero HTTP.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") +@pytest.mark.parametrize("flag", ["LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_OFFLINE_MODE"]) +async def test_refetch_respects_local_env_override(monkeypatch, flag): + """LITELLM_LOCAL_MODEL_COST_MAP=True or LITELLM_OFFLINE_MODE=True short-circuits to the bundled backup, zero HTTP.""" + monkeypatch.setenv(flag, "True") def _fail(request): raise AssertionError("no HTTP request should be made when local map is forced") @@ -708,9 +709,10 @@ def test_boot_load_falls_back_when_urlsplit_rejects_the_url(): 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") +@pytest.mark.parametrize("flag", ["LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_OFFLINE_MODE"]) +def test_boot_load_respects_local_env_override(monkeypatch, flag): + """LITELLM_LOCAL_MODEL_COST_MAP=True or LITELLM_OFFLINE_MODE=True still short-circuits to the backup with zero HTTP.""" + monkeypatch.setenv(flag, "True") def _fail(request): raise AssertionError("no HTTP request should be made when local map is forced") diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 79645ffe8f6..561ac9b2b33 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -42,6 +42,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( + BilledRoundsSpend, Logging, StandardLoggingPayloadSetup, _extract_response_obj_and_hidden_params, @@ -75,6 +76,7 @@ import inspect from typing import List, Optional from litellm.types.utils import LiteLLMCommonStrings + @pytest.fixture def logging_obj(): return LitellmLogging( @@ -345,8 +347,10 @@ async def test_mcp_native_structured_replacement_must_match_returned_content( ) logging_obj.dynamic_success_callbacks = [NativeReplacement()] returned = await logging_obj.async_post_mcp_tool_call_hook( - kwargs={"original_response": result}, response_obj=result, - start_time=datetime.datetime.now(), end_time=datetime.datetime.now(), + kwargs={"original_response": result}, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), ) assert returned is result assert result.content == [TextContent(type="text", text="native-safe" if same_content else "final-safe")] @@ -355,7 +359,6 @@ async def test_mcp_native_structured_replacement_must_match_returned_content( assert "SECRET-1234" not in result.model_dump_json() - @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["none", "wrapper", "exception"]) @pytest.mark.parametrize("structured", [False, True]) @@ -373,8 +376,10 @@ async def test_mcp_direct_content_edit_invalidates_stale_structured_data(logging ) logging_obj.dynamic_success_callbacks = [DirectRedactor()] returned = await logging_obj.async_post_mcp_tool_call_hook( - kwargs={"original_response": result}, response_obj=result, - start_time=datetime.datetime.now(), end_time=datetime.datetime.now(), + kwargs={"original_response": result}, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), ) assert returned is result assert result.content == [TextContent(type="text", text="[REDACTED]")] @@ -402,6 +407,32 @@ def test_get_combined_callback_list_preserves_insertion_order(logging_obj): ) == ["prometheus", "langfuse", "datadog", "otel", "s3", "gcs_bucket", "arize", "logfire"] +class _InputEventRecorder(CustomLogger): + def __init__(self, label: str, events: list[str]) -> None: + self._label = label + self.events = events + + def log_pre_api_call(self, model: object, messages: object, kwargs: object) -> None: + self.events.append(f"{self._label} pre") + + def log_post_api_call(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.events.append(f"{self._label} post") + + +def test_input_events_reach_the_replacement_and_not_the_global_callback_it_replaced( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final[list[str]] = [] + monkeypatch.setattr(litellm, "input_callback", [_InputEventRecorder("global", events)]) + logging_obj.replace_global_callback(_InputEventRecorder("replacement", events)) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + logging_obj.pre_call(input="", api_key="") + logging_obj.post_call(original_response="ok", input="", api_key="") + + assert events == ["replacement pre", "replacement post"] + + def test_get_masked_api_base(logging_obj): api_base = "https://api.openai.com/v1" masked_api_base = logging_obj._get_masked_api_base(api_base) @@ -1822,13 +1853,19 @@ def test_mark_logging_complete_flags_by_stream_and_event( "async_failure", "sync_failure", ) + attempt_marker: Final = datetime_unit_test(2026, 1, 1) logging_obj.stream = stream + logging_obj.model_call_details["current_attempt_api_call_start_time"] = attempt_marker logging_obj.mark_logging_complete(event_type=event) expected_to_run: Final = stream and event in ("async_success", "sync_success") assert logging_obj.should_run_logging(event_type=event) is expected_to_run assert all(logging_obj.should_run_logging(event_type=other_event) for other_event in events if other_event != event) + if event == "async_failure": + assert logging_obj.model_call_details["async_failure_logged_attempt"] is attempt_marker + else: + assert "async_failure_logged_attempt" not in logging_obj.model_call_details @pytest.mark.parametrize("stream", [False, True]) @@ -2014,7 +2051,9 @@ def test_langtrace_callback_exports_to_api_trace_with_x_api_key( ) -> None: """The exporter must post to Langtrace's complete /api/trace path with the key in x-api-key, without leaking it into the process-wide OTEL_EXPORTER_OTLP_TRACES_HEADERS.""" + import requests_mock from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + from opentelemetry.sdk.trace.export import SpanExportResult from litellm.integrations.opentelemetry import OpenTelemetry from litellm.litellm_core_utils import litellm_logging as logging_module @@ -2036,8 +2075,12 @@ def test_langtrace_callback_exports_to_api_trace_with_x_api_key( assert type(logger) is OpenTelemetry and logger.callback_name == "langtrace" exporter: Final = logger._get_span_processor().span_exporter assert isinstance(exporter, OTLPSpanExporter) - assert exporter._endpoint == expected_endpoint - assert exporter._headers == {"x-api-key": api_key} + with requests_mock.Mocker() as collector: + collector.post(expected_endpoint, status_code=200) + assert exporter.export(()) is SpanExportResult.SUCCESS + assert [(request.url, request.headers.get("x-api-key")) for request in collector.request_history] == [ + (expected_endpoint, api_key) + ] assert "OTEL_EXPORTER_OTLP_TRACES_HEADERS" not in os.environ finally: logging_module._in_memory_loggers.clear() @@ -2581,7 +2624,6 @@ async def test_completed_success_selects_one_implemented_hook( assert (thread_id == loop_thread) is (kind != "sync") - @pytest.mark.parametrize("for_async_request,split_registries", [(False, False), (True, False), (True, True)]) @pytest.mark.parametrize("private_first", [False, True]) @pytest.mark.parametrize("response_kind", ["responses", "completion"]) @@ -3273,8 +3315,13 @@ async def test_shadow_snapshot_stays_private_and_is_invalidated_before_logging_g class RecordingShadowLogger(ShadowEvalLogger): async def async_log_success_event( - self, kwargs: Mapping[str, object], response_obj: object, start_time: object, - end_time: object, *, guardrail_snapshot: GuardrailRequestSnapshot | None = None, + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: object, + end_time: object, + *, + guardrail_snapshot: GuardrailRequestSnapshot | None = None, ) -> None: shadow_snapshots.append(guardrail_snapshot) await super().async_log_success_event( @@ -3283,13 +3330,20 @@ async def test_shadow_snapshot_stays_private_and_is_invalidated_before_logging_g class RecordingLogger(CustomLogger): async def async_log_success_event( - self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object, + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: object, + end_time: object, ) -> None: other_payloads.append(kwargs) class LoggingGuardrail(CustomGuardrail): async def async_logging_hook( - self, kwargs: dict[str, object], result: object, call_type: str, + self, + kwargs: dict[str, object], + result: object, + call_type: str, ) -> tuple[dict[str, object], object]: hook_snapshots.append(logging_obj.shadow_eval_request_snapshot) if hook_mode == "raises": @@ -3301,26 +3355,36 @@ async def test_shadow_snapshot_stays_private_and_is_invalidated_before_logging_g "user_api_key_hash": "test-key", } snapshot: Final = GuardrailRequestSnapshot.capture( - {"messages": [{"role": "user", "content": "snapshot-only"}]}, metadata, + {"messages": [{"role": "user", "content": "snapshot-only"}]}, + metadata, ) assert snapshot is not None shadow: Final = RecordingShadowLogger(prisma_provider=no_prisma, jobs_cache=InMemoryCache()) guardrail: Final = LoggingGuardrail( - guardrail_name="late-mask", default_on=True, + guardrail_name="late-mask", + default_on=True, event_hook=GuardrailEventHooks.pre_call if hook_mode == "disabled" else GuardrailEventHooks.logging_only, ) monkeypatch.setattr(litellm, "_async_success_callback", []) logging_obj: Final = LitellmLogging( - model="test-model", messages=[], stream=stream, call_type="anthropic_messages", - start_time=datetime.datetime.now(), litellm_call_id="private-snapshot", function_id="private-snapshot", + model="test-model", + messages=[], + stream=stream, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id="private-snapshot", + function_id="private-snapshot", dynamic_async_success_callbacks=[shadow, RecordingLogger(), guardrail], ) logging_obj.update_messages([{"role": "user", "content": "logged input"}]) logging_obj.update_environment_variables(litellm_params={"metadata": metadata}, optional_params={}) logging_obj.shadow_eval_request_snapshot = snapshot payload: Final = { - "id": "private-snapshot", "call_type": "anthropic_messages", "metadata": metadata, - "model_group": "test-model", "model_parameters": {}, + "id": "private-snapshot", + "call_type": "anthropic_messages", + "metadata": metadata, + "model_group": "test-model", + "model_parameters": {}, } await logging_obj.async_success_handler(result=ModelResponse(), standard_logging_object=payload) @@ -3777,7 +3841,9 @@ def test_sentry_event_scrubber_initialization(monkeypatch): call_args = mock_init.call_args[1] assert call_args["send_default_pii"] is False assert call_args["event_scrubber"].recursive is True - assert {name.lower() for name in SENTRY_PII_DENYLIST} <= {name.lower() for name in call_args["event_scrubber"].denylist} + assert {name.lower() for name in SENTRY_PII_DENYLIST} <= { + name.lower() for name in call_args["event_scrubber"].denylist + } assert call_args["before_send"] is call_args["before_send_transaction"] @@ -3794,7 +3860,9 @@ def test_sentry_send_default_pii_opt_in(monkeypatch): call_args = mock_init.call_args[1] assert call_args["send_default_pii"] is True - assert not {name.lower() for name in SENTRY_PII_DENYLIST} & {name.lower() for name in call_args["event_scrubber"].denylist} + assert not {name.lower() for name in SENTRY_PII_DENYLIST} & { + name.lower() for name in call_args["event_scrubber"].denylist + } def test_get_masked_values(): @@ -5064,12 +5132,10 @@ def test_reported_cost_respects_configured_pricing_and_additional_charges( else {} ), } - } - } - ) - logging_obj.model_call_details.update( - {"additional_response_cost": additional_cost, "cache_hit": cache_hit} + }, + }, ) + logging_obj.model_call_details.update({"additional_response_cost": additional_cost, "cache_hit": cache_hit}) logging_obj._provider_reported_cost = reported_cost response: Final = ModelResponse( model=pricing_model, @@ -6811,6 +6877,26 @@ def test_failure_handler_zeroes_spend_without_recovered_usage(logging_obj): assert payload["total_tokens"] == 0 +def test_failure_after_earlier_rounds_bills_only_the_failed_round_when_it_priced_no_breakdown( + logging_obj: LitellmLogging, +) -> None: + earlier_breakdown: Final = CostBreakdown(total_cost=0.3, margin_fixed_amount=0.1) + logging_obj.cost_breakdown = earlier_breakdown + logging_obj.model_call_details["litellm_params"] = {} + + logging_obj.record_failure_after_earlier_rounds( + BilledRoundsSpend(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.3, earlier_breakdown), + Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), + 0.0, + ) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "anthropic", "claude"), "tb") + + row: Final = logging_obj.model_call_details["standard_logging_object"] + assert (row["prompt_tokens"], row["completion_tokens"]) == (7, 0) + assert row["response_cost"] == pytest.approx(0.0) + assert row.get("cost_breakdown") is None or row["cost_breakdown"]["total_cost"] == pytest.approx(0.0) + + def test_set_cost_breakdown_stores_reasoning_cost(): """reasoning_cost is stored only when positive, mirroring the cache-cost fields.""" from datetime import datetime @@ -6996,8 +7082,11 @@ async def test_background_interaction_completion_rebills_after_in_progress_succe class SyncRecorder(CustomLogger): def log_success_event( - self, kwargs: dict[str, object], response_obj: object, - start_time: dt.datetime, end_time: dt.datetime, + self, + kwargs: dict[str, object], + response_obj: object, + start_time: dt.datetime, + end_time: dt.datetime, ) -> None: events.put((response_obj, kwargs.get("response_cost"))) logged.set() @@ -7030,11 +7119,11 @@ async def test_background_interaction_completion_rebills_after_in_progress_succe assert logging_obj.model_call_details["response_cost"] > 0 assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 - assert events.qsize() == 2 assert events.get_nowait() == (in_progress, None) assert events.get_nowait() == (completed, logging_obj.model_call_details["response_cost"]) + @pytest.mark.asyncio async def test_background_interaction_completion_prices_the_settled_body_itself(): """ @@ -8068,7 +8157,9 @@ def test_debugging_log_with_json_logs_tolerates_missing_headers(logging_obj, mon logging_obj.litellm_request_debug = True with patch("litellm.litellm_core_utils.litellm_logging.verbose_logger.warning") as warning: - logging_obj._print_llm_call_debugging_log(api_base="https://api.openai.com/v1", headers=None, additional_args={}) + logging_obj._print_llm_call_debugging_log( + api_base="https://api.openai.com/v1", headers=None, additional_args={} + ) assert "https://api.openai.com/v1" in warning.call_args.kwargs["extra"]["api_base"] @@ -10319,8 +10410,13 @@ class TestAzurePTUSpilloverCost: assert payload is not None model_map_value: Final = payload["model_map_information"]["model_map_value"] assert model_map_value is not None - assert model_map_value["input_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["input_cost_per_token"] - assert model_map_value["output_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["output_cost_per_token"] + assert ( + model_map_value["input_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["input_cost_per_token"] + ) + assert ( + model_map_value["output_cost_per_token"] + == litellm.model_cost[self.SERVED_MODEL]["output_cost_per_token"] + ) assert response_cost == pytest.approx( 100 * model_map_value["input_cost_per_token"] + 50 * model_map_value["output_cost_per_token"] ) @@ -10435,8 +10531,6 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost() assert logging_obj.response_cost_calculator(result=assembled) == 0.0042 - - def test_response_cost_calculator_prices_terminal_responses_event_from_its_response(): logging_obj: Final = _responses_stream_logging_obj() inner_response: Final = ResponsesAPIResponse( @@ -10666,9 +10760,7 @@ def test_standard_logging_retries(): ] ) - with patch.object( - customHandler, "log_failure_event", new=MagicMock() - ) as mock_client: + with patch.object(customHandler, "log_failure_event", new=MagicMock()) as mock_client: try: router.completion( model="gpt-3.5-turbo", @@ -10680,19 +10772,10 @@ def test_standard_logging_retries(): pass assert mock_client.call_count == 2 + assert mock_client.call_args_list[0].kwargs["kwargs"]["standard_logging_object"]["trace_id"] is not None assert ( - mock_client.call_args_list[0].kwargs["kwargs"]["standard_logging_object"][ - "trace_id" - ] - is not None - ) - assert ( - mock_client.call_args_list[0].kwargs["kwargs"]["standard_logging_object"][ - "trace_id" - ] - == mock_client.call_args_list[1].kwargs["kwargs"][ - "standard_logging_object" - ]["trace_id"] + mock_client.call_args_list[0].kwargs["kwargs"]["standard_logging_object"]["trace_id"] + == mock_client.call_args_list[1].kwargs["kwargs"]["standard_logging_object"]["trace_id"] ) @@ -10959,9 +11042,7 @@ def test_signoz_dispatch_requires_an_endpoint(monkeypatch): custom_logger_init_args={}, ) assert created is None - assert not [ - cb for cb in logging_module._in_memory_loggers if getattr(cb, "callback_name", None) == "signoz" - ] + assert not [cb for cb in logging_module._in_memory_loggers if getattr(cb, "callback_name", None) == "signoz"] finally: logging_module._in_memory_loggers.clear() monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) @@ -11103,7 +11184,11 @@ def test_standard_logging_payload_reports_the_deployment_entry_the_request_was_b } } logging_obj.update_environment_variables( - model="gpt-5.4-nano", user="", optional_params={}, litellm_params=litellm_params, custom_llm_provider="openai" + model="gpt-5.4-nano", + user="", + optional_params={}, + litellm_params=litellm_params, + custom_llm_provider="openai", ) response: Final = ModelResponse( model="gpt-5.4-nano", @@ -11415,10 +11500,7 @@ def test_enterprise_alerting_loggers_falls_back_without_disabling_callback_contr ) logging_module._enterprise_alerting_loggers.cache_clear() for module_name in tuple(sys.modules): - if any( - module_name == prefix or module_name.startswith(f"{prefix}.") - for prefix in module_prefixes - ): + if any(module_name == prefix or module_name.startswith(f"{prefix}.") for prefix in module_prefixes): monkeypatch.delitem(sys.modules, module_name, raising=False) monkeypatch.setattr(sys, "meta_path", [BlockPagerDutyFinder(), *sys.meta_path]) @@ -11538,13 +11620,16 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest_asyncio.fixture(loop_scope="function") async def drain_logging_worker(isolate_litellm_state: None) -> AsyncIterator[None]: yield await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS: Final = LOGGING_WORKER_MAX_TIME_PER_COROUTINE + 5.0 + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -11582,6 +11667,7 @@ def isolate_litellm_state(): if attr in _DEFAULTS: setattr(litellm, attr, _DEFAULTS[attr]) + _LIST_ATTRS = ( "callbacks", "success_callback", @@ -11609,6 +11695,7 @@ _SCALAR_ATTRS = ( _DEFAULTS: dict = {} + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -11631,6 +11718,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "response_obj,expected_values", @@ -11678,6 +11766,7 @@ def test_get_usage(response_obj, expected_values): assert usage.completion_tokens == expected_values[1] assert usage.total_tokens == expected_values[2] + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_usage_from_image_generation_response(): """ @@ -11725,6 +11814,7 @@ def test_get_usage_from_image_generation_response(): assert usage.completion_tokens_details.image_tokens == 272 assert usage.completion_tokens_details.text_tokens == 100 + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_additional_headers(): additional_headers = { @@ -11765,10 +11855,12 @@ def test_get_additional_headers(): assert additional_logging_headers.get("llm_provider-request-id") == "req_01F6CycZZPSHKRCCctcS1Vto" assert additional_logging_headers.get("llm_provider-anthropic-ratelimit-requests-reset") == "2024-10-29T23:57:40Z" + def all_fields_present(standard_logging_metadata: StandardLoggingMetadata): for field in StandardLoggingMetadata.__annotations__.keys(): assert field in standard_logging_metadata + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "metadata_key, metadata_value", @@ -11800,6 +11892,7 @@ def test_get_standard_logging_metadata(metadata_key, metadata_value): # Assert that the specific metadata field is set correctly assert standard_logging_metadata[metadata_key] == metadata_value + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_user_api_key_hash(): valid_hash = "a" * 64 # 64 character string @@ -11807,6 +11900,7 @@ def test_get_standard_logging_metadata_user_api_key_hash(): result = StandardLoggingPayloadSetup.get_standard_logging_metadata(metadata) assert result["user_api_key_hash"] == valid_hash + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_invalid_user_api_key(): invalid_hash = "not_a_valid_hash" @@ -11815,6 +11909,7 @@ def test_get_standard_logging_metadata_invalid_user_api_key(): all_fields_present(result) assert result["user_api_key_hash"] is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_non_string_user_api_key(): """Non-string user_api_key should not be set as user_api_key_hash.""" @@ -11823,6 +11918,7 @@ def test_get_standard_logging_metadata_non_string_user_api_key(): all_fields_present(result) assert result["user_api_key_hash"] is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_none_user_api_key(): """None user_api_key should not be set as user_api_key_hash.""" @@ -11831,6 +11927,7 @@ def test_get_standard_logging_metadata_none_user_api_key(): all_fields_present(result) assert result["user_api_key_hash"] is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_invalid_keys(): metadata = { @@ -11844,6 +11941,7 @@ def test_get_standard_logging_metadata_invalid_keys(): assert "invalid_key" not in result assert "another_invalid_key" not in result + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_cleanup_timestamps(): """Test cleanup_timestamps with different input types""" @@ -11878,6 +11976,7 @@ def test_cleanup_timestamps(): with pytest.raises(ValueError, match="start_time is required, got=invalid of type "): StandardLoggingPayloadSetup.cleanup_timestamps("invalid", end_float, completion_float) + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_model_cost_information(): """Test get_model_cost_information with different inputs""" @@ -11910,6 +12009,7 @@ def test_get_model_cost_information(): # assert all fields in StandardLoggingModelInformation are present assert all(field in result for field in StandardLoggingModelInformation.__annotations__) + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_model_cost_information_custom_pricing_uses_base_model(): result = StandardLoggingPayloadSetup.get_model_cost_information( @@ -11921,6 +12021,7 @@ def test_get_model_cost_information_custom_pricing_uses_base_model(): assert result["model_map_value"] is not None assert result["model_map_key"] != "invoke_test_claude" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_standard_logging_payload_uses_deployment_when_no_base_model(): """metadata["deployment"] is used for cost-map lookup when base_model is not set.""" @@ -11978,6 +12079,7 @@ def test_standard_logging_payload_uses_deployment_when_no_base_model(): assert payload["model_map_information"]["model_map_value"] is not None assert payload["model_map_information"]["model_map_key"] != "invoke_test_claude" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_hidden_params(): """Test get_hidden_params with different inputs""" @@ -12013,6 +12115,7 @@ def test_get_hidden_params(): # assert all fields in StandardLoggingHiddenParams are present assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_final_response_obj(): """Test get_final_response_obj with different input types and redaction scenarios""" @@ -12044,6 +12147,7 @@ def test_get_final_response_obj(): # Reset litellm.turn_off_message_logging to its original value litellm.turn_off_message_logging = False + def testget_standard_logging_payload_trace_id(): """Test get_standard_logging_payload_trace_id with different input scenarios""" # Test case 1: When litellm_trace_id is provided in litellm_params @@ -12081,6 +12185,7 @@ def testget_standard_logging_payload_trace_id(): assert result == "12345" assert isinstance(result, str) + def testget_standard_logging_payload_trace_id_prioritizes_trace_id_when_flag_on(monkeypatch): """With request_correlation_in_logs on, an explicit litellm_trace_id wins over litellm_session_id.""" from unittest.mock import MagicMock @@ -12095,6 +12200,7 @@ def testget_standard_logging_payload_trace_id_prioritizes_trace_id_when_flag_on( ) assert result == "the-trace-id" + def testget_standard_logging_payload_trace_id_prioritizes_session_id_when_flag_off(monkeypatch): """With request_correlation_in_logs off (default), legacy behavior is preserved: litellm_session_id still wins over litellm_trace_id.""" @@ -12110,6 +12216,7 @@ def testget_standard_logging_payload_trace_id_prioritizes_session_id_when_flag_o ) assert result == "the-session-id" + def testget_standard_logging_payload_session_id_when_flag_on(monkeypatch): """Test get_standard_logging_payload_session_id with different input scenarios, flag enabled""" from unittest.mock import MagicMock @@ -12162,6 +12269,7 @@ def testget_standard_logging_payload_session_id_when_flag_on(monkeypatch): ) assert result == "" + def testget_standard_logging_payload_session_id_empty_when_flag_off(monkeypatch): """When request_correlation_in_logs is off (default), session_id is always empty, even if litellm_session_id was explicitly supplied - preserves the pre-existing @@ -12178,6 +12286,7 @@ def testget_standard_logging_payload_session_id_empty_when_flag_off(monkeypatch) ) assert result == "" + @pytest.fixture def restore_verbose_logger_level() -> Iterator[None]: original_level: Final = verbose_logger.level @@ -12220,6 +12329,7 @@ def test_truncate_standard_logging_payload(): assert len(str(truncated["response"])) < 10_500 assert len(str(truncated["error_str"])) < 10_500 + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_truncate_standard_logging_payload_keeps_a_partial_payload_intact(): """A payload built with only some of its fields comes back with exactly those keys and values""" @@ -12228,12 +12338,14 @@ def test_truncate_standard_logging_payload_keeps_a_partial_payload_intact(): assert _custom_logger.truncate_standard_logging_payload_content(partial_payload) == partial_payload + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_strip_trailing_slash(): common_api_base = "https://api.test.com" assert StandardLoggingPayloadSetup.strip_trailing_slash(common_api_base + "/") == common_api_base assert StandardLoggingPayloadSetup.strip_trailing_slash(common_api_base) == common_api_base + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_error_information(): """Test get_error_information with different types of exceptions""" @@ -12270,6 +12382,7 @@ def test_get_error_information(): assert result["llm_provider"] == "openai" assert result["error_message"] == "litellm.RateLimitError: Test error" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_response_time(): """Test get_response_time with different streaming scenarios""" @@ -12305,6 +12418,7 @@ def test_get_response_time(): # For streaming, should return completion_start_time - start_time assert response_time == 2.0 + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "metadata, expected_requester_metadata", @@ -12327,6 +12441,7 @@ def test_standard_logging_metadata_requester_metadata(metadata, expected_request result = StandardLoggingPayloadSetup.get_standard_logging_metadata(metadata) assert result["requester_metadata"] == expected_requester_metadata + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_cost_breakdown_in_standard_logging_payload(): """ @@ -12416,6 +12531,7 @@ def test_cost_breakdown_in_standard_logging_payload(): print("✅ Cost breakdown test passed!") + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_cost_breakdown_missing_in_standard_logging_payload(): """ @@ -12476,6 +12592,7 @@ def test_cost_breakdown_missing_in_standard_logging_payload(): print("✅ Cost breakdown missing test passed!") + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "use_combined_usage_object", @@ -12563,6 +12680,7 @@ def test_usage_dict_roundtrip_in_payload(use_combined_usage_object): assert usage_obj["completion_tokens"] == 58 assert usage_obj["total_tokens"] == 100 + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_standard_logging_payload_uses_actual_model_for_azure_router(): from litellm.litellm_core_utils.litellm_logging import ( @@ -12611,6 +12729,7 @@ def test_standard_logging_payload_uses_actual_model_for_azure_router(): assert payload is not None assert payload["model"] == "azure_ai/gpt-5-nano-2025-08-07" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_standard_logging_payload_uses_actual_model_for_azure_router_with_underscore(): from litellm.litellm_core_utils.litellm_logging import ( @@ -12659,6 +12778,7 @@ def test_standard_logging_payload_uses_actual_model_for_azure_router_with_unders assert payload is not None assert payload["model"] == "azure_ai/gpt-5-nano-2025-08-07" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_basic(): """ @@ -12690,6 +12810,7 @@ def test_merge_litellm_metadata_basic(): assert result["model_info"] == {"id": "model-123"} assert result["tags"] == ["tag1", "tag2"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_precedence(): """ @@ -12716,6 +12837,7 @@ def test_merge_litellm_metadata_precedence(): # litellm_metadata values should only be included if not in metadata assert result["model_group"] == "gpt-4-group" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_skip_non_serializable(): """ @@ -12748,6 +12870,7 @@ def test_merge_litellm_metadata_skip_non_serializable(): assert result["safe_field"] == "safe_value" assert result["model_group"] == "gpt-4-group" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_empty_params(): """ @@ -12783,6 +12906,7 @@ def test_merge_litellm_metadata_empty_params(): result = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) assert result == {} + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_bedrock_passthrough_scenario(): """ @@ -12842,8 +12966,10 @@ def test_merge_litellm_metadata_bedrock_passthrough_scenario(): # Verify total number of fields (9 user fields + 4 model fields = 13) assert len(result) == 13 + service_logger = ServiceLogging() + def setup_logging(): return Logging( model="gpt-5.5", @@ -12933,6 +13059,7 @@ def test_get_callback_name(): # Test string callback assert logging._get_callback_name("callback_string") == "callback_string" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_is_internal_litellm_proxy_callback(): """ @@ -12953,6 +13080,7 @@ def test_is_internal_litellm_proxy_callback(): # Test string callback assert logging._is_internal_litellm_proxy_callback("callback_string") == False + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_should_run_sync_callbacks_for_async_calls(): """ @@ -12978,6 +13106,7 @@ def test_should_run_sync_callbacks_for_async_calls(): litellm.success_callback = [PROXY_MaxIterationsHandler] assert logging._should_run_sync_callbacks_for_async_calls() == False + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_remove_internal_litellm_callbacks(): logging = setup_logging() @@ -12999,6 +13128,7 @@ def test_remove_internal_litellm_callbacks(): assert PROXY_MaxIterationsHandler not in filtered assert PROXY_CacheControlCheck not in filtered + @pytest.mark.asyncio async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch): """ @@ -13101,14 +13231,9 @@ class CompletionCustomHandler( assert metadata_value is None or isinstance(metadata_value, dict) if metadata_value is not None: if litellm.turn_off_message_logging is True: - assert ( - metadata_value["raw_request"] - is LiteLLMCommonStrings.redacted_by_litellm.value - ) + assert metadata_value["raw_request"] is LiteLLMCommonStrings.redacted_by_litellm.value else: - assert "raw_request" not in metadata_value or isinstance( - metadata_value["raw_request"], str - ) + assert "raw_request" not in metadata_value or isinstance(metadata_value["raw_request"], str) except Exception: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) @@ -13157,23 +13282,18 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponseStream) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -13206,9 +13326,7 @@ class CompletionCustomHandler( ) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["litellm_params"]["api_base"], str) @@ -13218,10 +13336,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["user"], (str, type(None))) assert ( isinstance(kwargs["input"], list) - and ( - isinstance(kwargs["input"][0], dict) - or isinstance(kwargs["input"][0], str) - ) + and (isinstance(kwargs["input"][0], dict) or isinstance(kwargs["input"][0], str)) ) or isinstance(kwargs["input"], (dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert isinstance( @@ -13249,9 +13364,7 @@ class CompletionCustomHandler( assert response_obj == None ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) @@ -13259,15 +13372,12 @@ class CompletionCustomHandler( assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None ) assert isinstance(kwargs["additional_args"], (dict, type(None))) @@ -13285,9 +13395,7 @@ class CompletionCustomHandler( assert isinstance(messages, list) and isinstance(messages[0], dict) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) @@ -13299,9 +13407,7 @@ class CompletionCustomHandler( async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - print( - "in async_log_success_event", kwargs, response_obj, start_time, end_time - ) + print("in async_log_success_event", kwargs, response_obj, start_time, end_time) self.states.append("async_success") ## START TIME assert isinstance(start_time, datetime) @@ -13330,9 +13436,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -13364,9 +13468,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, str, dict)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) or kwargs["original_response"] == None @@ -13507,3 +13609,36 @@ def test_legacy_s3_logs_a_streamed_anthropic_messages_call_handed_off_from_the_a assert s3_writer.log_event.call_count == expected_s3_writes if expected_s3_writes: assert s3_writer.log_event.call_args.kwargs["response_obj"] is complete_response + + +_GO_PRICED_MESSAGE: Final = { + "id": "msg_go", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1200, "output_tokens": 300, "cost": 0, "opencode": {"list_cost": 0.00412}}, +} + + +@pytest.mark.parametrize( + ("usage", "expected_header"), + ( + ({"input_tokens": 1200, "output_tokens": 300, "cost": 0, "opencode": {"list_cost": 0.00412}}, 0.00412), + ({"input_tokens": 1200, "output_tokens": 300}, None), + ), + ids=("list-cost-reported", "no-list-cost-keeps-the-price-map"), +) +def test_native_messages_body_carries_the_reported_list_cost_into_the_priced_response(usage, expected_header): + logging_obj = _anthropic_messages_logging_obj() + logging_obj.stream = False + logging_obj.optional_params = {} + logging_obj.model_call_details["custom_llm_provider"] = "opencode_go" + + priced = logging_obj.billable_result({**_GO_PRICED_MESSAGE, "usage": usage}) + + assert isinstance(priced, ModelResponse) + headers = priced._hidden_params.get("additional_headers") or {} + assert headers.get("llm_provider-x-litellm-response-cost") == expected_header diff --git a/tests/unit/litellm_core_utils/test_logging_callback_manager.py b/tests/unit/litellm_core_utils/test_logging_callback_manager.py index f8260702068..25a52137d65 100644 --- a/tests/unit/litellm_core_utils/test_logging_callback_manager.py +++ b/tests/unit/litellm_core_utils/test_logging_callback_manager.py @@ -145,6 +145,7 @@ def test_remove_callback_from_all_lists(): obj = TestLogger() manager.add_litellm_callback(obj) + manager.add_litellm_input_callback(obj) manager.add_litellm_success_callback(obj) manager.add_litellm_failure_callback(obj) manager.add_litellm_async_success_callback(obj) @@ -153,6 +154,7 @@ def test_remove_callback_from_all_lists(): manager.remove_callback_from_all_lists(obj) assert obj not in litellm.callbacks + assert obj not in litellm.input_callback assert obj not in litellm.success_callback assert obj not in litellm.failure_callback assert obj not in litellm._async_success_callback diff --git a/tests/unit/litellm_core_utils/test_max_streaming_duration.py b/tests/unit/litellm_core_utils/test_max_streaming_duration.py index fc8daab3899..234bc83e3c3 100644 --- a/tests/unit/litellm_core_utils/test_max_streaming_duration.py +++ b/tests/unit/litellm_core_utils/test_max_streaming_duration.py @@ -67,12 +67,13 @@ class TestCustomStreamWrapperMaxDuration: @pytest.mark.asyncio async def test_should_raise_on_async_anext_when_exceeded(self): - """__anext__ should check the limit before iterating, dispatching the - same failure-callback/logging path every other stream failure goes - through (dispatch_failure_handlers is async on the real Logging class, - so the mock needs to be awaitable too).""" + """__anext__ checks the limit before iterating and schedules failure logging.""" wrapper = _make_custom_stream_wrapper() - wrapper.logging_obj.dispatch_failure_handlers = AsyncMock() + + async def _noop() -> None: + return None + + wrapper.logging_obj.dispatch_attempt_failure_handlers = MagicMock(side_effect=lambda *_: _noop()) wrapper._stream_created_time = time.time() - 20 with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): with pytest.raises(litellm.Timeout): diff --git a/tests/unit/litellm_core_utils/test_offline_mode.py b/tests/unit/litellm_core_utils/test_offline_mode.py new file mode 100644 index 00000000000..f5a4786837d --- /dev/null +++ b/tests/unit/litellm_core_utils/test_offline_mode.py @@ -0,0 +1,17 @@ +import pytest + +from litellm.litellm_core_utils.offline_mode import is_offline_mode + + +@pytest.mark.parametrize( + ("value", "expected"), + [("true", True), ("True", True), (" TRUE ", True), ("false", False), ("", False), ("1", False), ("yes", False)], +) +def test_is_offline_mode_reads_true_case_insensitively(monkeypatch, value, expected): + monkeypatch.setenv("LITELLM_OFFLINE_MODE", value) + assert is_offline_mode() is expected + + +def test_is_offline_mode_is_off_when_unset(monkeypatch): + monkeypatch.delenv("LITELLM_OFFLINE_MODE", raising=False) + assert is_offline_mode() is False diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index a91888fea79..6a734ca3b70 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -57,6 +57,8 @@ def _provider_fields() -> dict[str, object]: "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], "citations": [{"cited_text": SECRET}], "thought_signatures": [SIGNATURE], + "compaction_blocks": [{"type": "compaction", "content": SECRET}], + "compaction_delta": {"type": "compaction_delta", "content": SECRET}, } @@ -165,6 +167,26 @@ def test_streaming_redaction_redacts_delta_provider_fields_without_touching_the_ assert fields == _provider_fields() +@pytest.mark.parametrize("shape", ["object", "dict"]) +def test_redaction_hides_a_compaction_summary_in_responses_api_output(shape: str) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[ + {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": SECRET})} + ], + ) + result: Final = response if shape == "object" else response.model_dump() + + redacted: Final = perform_redaction({}, result) + + assert SECRET not in str(redacted), redacted + assert SECRET in str(result) + + @pytest.mark.parametrize("surface", ("typed", "dict", "standard", "callback", "helper")) def test_responses_redaction_removes_instructions_without_changing_the_response(surface: str) -> None: response: Final = litellm.ResponsesAPIResponse.model_validate( diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 68aabc3c01e..cc1fa1b8803 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,6 +1,7 @@ import asyncio import datetime import json +import queue import sys import threading import time @@ -13,8 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest from openai.types.chat import ChatCompletionChunk -from pydantic import BaseModel - +from pydantic import BaseModel, TypeAdapter import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var @@ -36,6 +36,7 @@ from litellm.types.utils import ( ModelResponseStream, PromptTokensDetailsWrapper, StreamingChoices, + StandardLoggingPayload, Usage, ) from litellm.utils import ModelResponseListIterator @@ -603,6 +604,91 @@ async def test_streaming_handler_with_usage( assert chunk_has_usage +def _usage_test_logging( + caller_stream_options: dict[str, bool] | None, stream_usage_injected: bool = False +) -> Logging: + logging_obj: Final = Logging( + model="hosted_vllm/m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=1_790_000_000.0, + litellm_call_id="12345", + function_id="1245", + kwargs={"_litellm_stream_usage_injected": True} if stream_usage_injected else {}, + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="hosted_vllm/m", stream_options=caller_stream_options + ) + return logging_obj + + +@pytest.mark.parametrize( + ("caller_stream_options", "stream_usage_injected", "usage_chunk_choice_count"), + [ + ({"include_usage": True}, False, 0), + (None, False, 1), + ({"include_usage": True}, True, 1), + ], + ids=["caller-asked", "provider-handler-injected", "proxy-injected"], +) +@pytest.mark.parametrize("plain_iterator", [False, True]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_usage_chunk_has_empty_choices_only_for_callers_that_asked( + sync_mode: bool, + plain_iterator: bool, + caller_stream_options: dict[str, bool] | None, + stream_usage_injected: bool, + usage_chunk_choice_count: int, +): + usage: Final = Usage(prompt_tokens=5, completion_tokens=1, total_tokens=6) + provider_chunks: Final = [ + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="ok"))]), + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), + ModelResponseStream(id="c1", choices=[], provider_specific_fields={"citations": ["a"]}), + ModelResponseStream(id="c1", choices=[], usage=usage), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=( + iter(provider_chunks) if plain_iterator else ModelResponseListIterator(model_responses=provider_chunks) + ), + model="hosted_vllm/m", + custom_llm_provider="hosted_vllm", + logging_obj=_usage_test_logging(caller_stream_options, stream_usage_injected), + stream_options={"include_usage": True}, + ) + + chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + + assert [len(chunk.choices) for chunk in chunks] == [1, 1, 0, usage_chunk_choice_count] + assert chunks[2].provider_specific_fields == {"citations": ["a"]} + assert [getattr(chunk, "usage", None) is not None for chunk in chunks] == [False, False, False, True] + assert chunks[-1].usage.total_tokens == 6 + + +@pytest.mark.asyncio +async def test_plain_iterator_stream_keeps_usage_on_its_chunk_when_usage_was_not_requested(): + usage: Final = Usage(prompt_tokens=5, completion_tokens=1, total_tokens=6) + provider_chunks: Final = [ + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="ok"))]), + ModelResponseStream( + id="c1", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], usage=usage + ), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(provider_chunks), + model="hosted_vllm/m", + custom_llm_provider="hosted_vllm", + logging_obj=_usage_test_logging(None), + ) + + chunks: Final = [chunk async for chunk in wrapper] + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"] + assert chunks[-1].usage.total_tokens == 6 + + @pytest.mark.parametrize("sync_mode", [False]) @pytest.mark.asyncio @pytest.mark.flaky(reruns=3) @@ -2734,6 +2820,123 @@ def test_raise_on_model_repetition_tolerates_empty_choices( wrapper.raise_on_model_repetition() +_STUCK_CHUNK: Final = "the model is stuck on this" +_STUCK_PROVIDER_CHUNK_CAP: Final = 10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT + + +class _StuckProvider: + def __init__(self) -> None: + self.closed_early = threading.Event() + self.chunks_sent = 0 + + def _next_chunk(self) -> ModelResponseStream: + self.chunks_sent += 1 + return _make_chunk(_STUCK_CHUNK) + + def sync_stream(self) -> Iterator[ModelResponseStream]: + try: + yield from (self._next_chunk() for _ in range(_STUCK_PROVIDER_CHUNK_CAP)) + except GeneratorExit: + self.closed_early.set() + raise + + async def async_stream(self) -> AsyncIterator[ModelResponseStream]: + try: + for _ in range(_STUCK_PROVIDER_CHUNK_CAP): + yield self._next_chunk() + except GeneratorExit: + self.closed_early.set() + raise + + +class _StuckStreamFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: Final = queue.SimpleQueue[str]() + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"])) + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"])) + + +def _stuck_stream_wrapper( + provider_stream: Iterator[ModelResponseStream] | AsyncIterator[ModelResponseStream], + call_type: Literal["completion", "acompletion"], + recorder: _StuckStreamFailureRecorder, +) -> CustomStreamWrapper: + logging_obj: Final = Logging( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type=call_type, + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id=f"stuck-stream-{call_type}", + function_id=f"stuck-stream-{call_type}", + dynamic_failure_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="test-model", + optional_params={}, + litellm_params=get_litellm_params(custom_llm_provider="openai", acompletion=call_type == "acompletion"), + ) + return CustomStreamWrapper( + completion_stream=provider_stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + +@pytest.mark.asyncio +async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): + """https://github.com/BerriAI/litellm/issues/19428 and https://github.com/BerriAI/litellm/issues/13786""" + provider: Final = _StuckProvider() + recorder: Final = _StuckStreamFailureRecorder() + wrapper: Final = _stuck_stream_wrapper(provider.async_stream(), "acompletion", recorder) + + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in wrapper: + pass + assert provider.closed_early.is_set() + assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1 + chunks_sent_at_failure: Final = provider.chunks_sent + + with pytest.raises(StopAsyncIteration): + await wrapper.__anext__() + await wrapper.wait_for_failure_logging(timeout=5) + + assert provider.chunks_sent == chunks_sent_at_failure + assert "repeating the same chunk" in recorder.failures.get_nowait() + assert recorder.failures.empty() + + +def test_sync_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): + """Sync counterpart of the async test above""" + provider: Final = _StuckProvider() + recorder: Final = _StuckStreamFailureRecorder() + wrapper: Final = _stuck_stream_wrapper(provider.sync_stream(), "completion", recorder) + + with pytest.raises(Exception, match="repeating the same chunk"): + for _ in wrapper: + pass + assert provider.closed_early.is_set() + assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1 + chunks_sent_at_failure: Final = provider.chunks_sent + assert "repeating the same chunk" in recorder.failures.get(timeout=5) + + with pytest.raises(StopIteration): + next(wrapper) + + assert provider.chunks_sent == chunks_sent_at_failure + with pytest.raises(queue.Empty): + recorder.failures.get(timeout=0.5) + + def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): """ Test that provider-reported usage from a post-finish_reason chunk @@ -4180,6 +4383,68 @@ def test_record_partial_usage_for_failure_stashes_usage_and_cost(): assert isinstance(logging_obj.model_call_details["response_cost"], float) +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: tuple[tuple[int, int], ...] = () + self.logged: asyncio.Event = asyncio.Event() + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + self.failures = (*self.failures, (payload["prompt_tokens"], payload["completion_tokens"])) + self.logged.set() + + +@pytest.mark.asyncio +async def test_log_stream_failure_bills_the_dropped_attempt_even_after_the_router_discards_its_usage() -> None: + """The dropped attempt's failure log is queued, and the router discards the stashed usage as it opens the + fallback hop before that log runs. The log must still bill what the dropped attempt consumed, and the shared + logging object must not be marked as already failed, since the hop's own failure still has to reach the + integrations.""" + recorder: Final = _FailureRecorder() + logging_obj: Final = Logging( + model="gpt-5.6", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="partial-usage-copy", + function_id="1245", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-5.6", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-partial-copy", + created=1742056047, + model="gpt-5.6", + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="The", role="assistant"))], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + + with pytest.raises(MidStreamFallbackError): + wrapper._log_stream_failure_and_raise(httpx.ReadError("connection dropped")) + logging_obj.carry_failed_attempt_usage() + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + + assert recorder.failures == ((30, 1),) + assert logging_obj.model_call_details.get("has_logged_async_failure") is not True + + def test_record_partial_usage_for_failure_noop_without_chunks(): """With no chunks delivered there is nothing billed to recover, so the failure stash must stay absent and not force a zero-usage row. diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 309cbd33131..88af656ad2e 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -14,7 +14,7 @@ from concurrent.futures import Future, ThreadPoolExecutor, wait from itertools import accumulate, chain from pathlib import Path from typing import Final -from unittest.mock import MagicMock +from unittest.mock import MagicMock, sentinel import anyio.to_thread import pytest @@ -27,7 +27,7 @@ import litellm from litellm import decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +from litellm.constants import HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, TOKEN_COUNTER_MAX_CONCURRENT_COUNTS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( _encoding_count, @@ -940,10 +940,17 @@ def test_token_counter(): import unittest -from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding +from litellm.utils import ( + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _load_huggingface_tokenizer, + _select_tokenizer_helper, + claude_json_str, + encoding, +) # Clear the cache at module load to ensure clean state _load_huggingface_tokenizer.cache_clear() +_HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() class TestTokenizerSelection(unittest.TestCase): @@ -955,6 +962,7 @@ class TestTokenizerSelection(unittest.TestCase): --dist=loadscope (tests from same file run on same worker). """ _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() @patch("litellm.utils.tokenizer_dispatch.from_pretrained") def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): @@ -1038,6 +1046,69 @@ class TestTokenizerSelection(unittest.TestCase): finally: monkeypatch.undo() + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") + def test_failed_huggingface_tokenizer_download_is_not_retried_on_every_count( + self, mock_from_pretrained: MagicMock + ): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + try: + with patch("litellm.utils.time.monotonic", return_value=0.0): + mock_from_pretrained.side_effect = RuntimeError("HF unavailable") + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + text: Final = "Count the tokens for this message." + expected: Final = len(encoding.encode(text, disallowed_special=())) + + first_count: Final = litellm.token_counter(model=model, text=text) + first_attempts: Final = mock_from_pretrained.call_count + second_count: Final = litellm.token_counter(model=model, text=text) + + assert first_attempts > 0 + assert mock_from_pretrained.call_count == first_attempts + assert first_count == expected + assert second_count == expected + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + + @patch("litellm.utils.time.monotonic", return_value=299.0) + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") + def test_failed_huggingface_tokenizer_download_is_retried_after_retry_window( + self, mock_from_pretrained: MagicMock, mock_monotonic: MagicMock + ): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + try: + def from_pretrained_side_effect(_: str) -> object: + if mock_from_pretrained.call_count == 1: + mock_monotonic.return_value = 322.0 + raise RuntimeError("HF unavailable") + return sentinel.tokenizer + + mock_from_pretrained.side_effect = from_pretrained_side_effect + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + + first_result: Final = _select_tokenizer_helper(model) + assert first_result["type"] == "openai_tokenizer" + mock_monotonic.return_value = 323.0 + second_result: Final = _select_tokenizer_helper(model) + assert second_result["type"] == "openai_tokenizer", "tokenizer was retried at t=323" + assert mock_from_pretrained.call_count == 1 + + mock_monotonic.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS - 1 + third_result: Final = _select_tokenizer_helper(model) + assert third_result["type"] == "openai_tokenizer" + assert mock_from_pretrained.call_count == 1 + + mock_monotonic.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1 + retried_result: Final = _select_tokenizer_helper(model) + assert retried_result["type"] == "huggingface_tokenizer" + assert retried_result["tokenizer"] is sentinel.tokenizer + assert mock_from_pretrained.call_count == 2 + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _CODE_SAMPLE: Final = " for item in items:\n total += item.price * item.quantity\n" * 20 @@ -1102,8 +1173,30 @@ def test_disable_hf_tokenizer_download_keeps_deepseek_v4_on_cl100k(monkeypatch: assert counted == litellm.token_counter(model="", text=_CODE_SAMPLE) +@pytest.mark.parametrize("model", ["llama-2-7b", "meta-llama/llama-3-8b-instruct"]) +def test_offline_mode_counts_downloaded_tokenizer_models_with_tiktoken(monkeypatch: pytest.MonkeyPatch, model: str): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") + + with patch("litellm.utils.tokenizer_dispatch.from_pretrained") as from_pretrained: + result: Final = _select_tokenizer_helper(model) + + from_pretrained.assert_not_called() + assert result["type"] == "openai_tokenizer" + + +def test_offline_mode_keeps_deepseek_v4_on_its_bundled_tokenizer(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") + + assert litellm.token_counter(model="deepseek/deepseek-v4-flash", text=_CODE_SAMPLE) == _deepseek_token_count( + _CODE_SAMPLE + ) + + def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() start: Final = threading.Barrier(8) def first_tokenizer(_: int) -> object: @@ -1116,6 +1209,44 @@ def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): assert len({id(tokenizer) for tokenizer in tokenizers}) == 1 +def test_remembered_tokenizer_failure_warns_once_instead_of_on_every_count(caplog: pytest.LogCaptureFixture): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + try: + with ( + patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable")), + caplog.at_level("WARNING", logger="LiteLLM"), + ): + types: Final = tuple(_select_tokenizer_helper(model)["type"] for _ in range(3)) + fallback_warnings: Final = [r for r in caplog.records if "token counts may be approximate" in r.getMessage()] + assert types == ("openai_tokenizer",) * 3 + assert len(fallback_warnings) == 1 + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + + +def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered(): + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted + + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + forked: Final = ForkedAfterNativeRuntimeStarted("native runtime started before fork") + try: + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=forked): + with pytest.raises(ForkedAfterNativeRuntimeStarted): + _select_tokenizer_helper(model) + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", return_value=sentinel.tokenizer): + loaded: Final = _select_tokenizer_helper(model) + assert loaded["type"] == "huggingface_tokenizer" + assert loaded["tokenizer"] is sentinel.tokenizer + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + + def test_token_counter_with_anthropic_tool_use(): """ Test that _count_anthropic_content() correctly handles tool_use blocks. diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 354c91920fc..61e797d7166 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -406,16 +406,26 @@ def test_huggingface_encoding_exposes_the_tokenizers_lookup_and_mutation_surface @pytest.mark.parametrize("log_level", ["WARNING", "ERROR"]) def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeypatch, log_level): from unittest.mock import patch - from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper + + from litellm.utils import ( + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _load_huggingface_tokenizer, + _select_tokenizer_helper, + ) monkeypatch.setenv("LITELLM_RUST", "false") _load_huggingface_tokenizer.cache_clear() - with caplog.at_level(log_level, logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": None}): - result = _select_tokenizer_helper("llama-2") - assert result["type"] == "openai_tokenizer" - assert result["tokenizer"].encode("hello") - assert ("token counts may be approximate" in caplog.text) is (log_level == "WARNING") - assert ("install tokenizers" in caplog.text) is (log_level == "WARNING") + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + try: + with caplog.at_level(log_level, logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": None}): + result = _select_tokenizer_helper("llama-2") + assert result["type"] == "openai_tokenizer" + assert result["tokenizer"].encode("hello") + assert ("token counts may be approximate" in caplog.text) is (log_level == "WARNING") + assert ("install tokenizers" in caplog.text) is (log_level == "WARNING") + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) @@ -471,10 +481,16 @@ def test_added_token_return_annotation_resolves_without_optional_import(python_i def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): from types import SimpleNamespace from unittest.mock import patch - from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper + + from litellm.utils import ( + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _load_huggingface_tokenizer, + _select_tokenizer_helper, + ) monkeypatch.setenv("LITELLM_RUST", "false") _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() model = "llama-2\r\nforged-model\x1b[31m\u2028\u2029" secret = "sk-" + "x" * 48 failure = OSError("download failed\r\nforged-error\x1b[31m api_key=" + secret) @@ -483,18 +499,26 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): raise failure dependency = SimpleNamespace(Tokenizer=SimpleNamespace(from_pretrained=fail_download)) - with caplog.at_level("WARNING", logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": dependency}): - result = _select_tokenizer_helper(model) - assert result["type"] == "openai_tokenizer" - assert result["tokenizer"].encode("hello") - message = next(record.getMessage() for record in caplog.records if "token counts may be approximate" in record.getMessage()) - assert "llama-2" in message - assert "download failed" in message - assert "forged-model" in message and "forged-error" in message - assert message.isascii() and message.isprintable() - assert secret not in message - assert "REDACTED" in message - assert "install tokenizers and huggingface-hub" in message + try: + with caplog.at_level("WARNING", logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": dependency}): + result = _select_tokenizer_helper(model) + assert result["type"] == "openai_tokenizer" + assert result["tokenizer"].encode("hello") + message = next( + record.getMessage() + for record in caplog.records + if "token counts may be approximate" in record.getMessage() + ) + assert "llama-2" in message + assert "download failed" in message + assert "forged-model" in message and "forged-error" in message + assert message.isascii() and message.isprintable() + assert secret not in message + assert "REDACTED" in message + assert "install tokenizers and huggingface-hub" in message + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) diff --git a/tests/unit/litellm_core_utils/test_url_utils.py b/tests/unit/litellm_core_utils/test_url_utils.py index 46bb91959c0..9b29eac208e 100644 --- a/tests/unit/litellm_core_utils/test_url_utils.py +++ b/tests/unit/litellm_core_utils/test_url_utils.py @@ -1,6 +1,8 @@ import asyncio import socket import threading +from collections.abc import Callable +from urllib.parse import urlsplit import httpx import pytest @@ -13,6 +15,8 @@ from litellm.litellm_core_utils.url_utils import ( assert_same_origin, encode_url_path_segment, encode_url_path_segments, + ensure_api_route, + join_api_base, validate_url, ) @@ -571,3 +575,80 @@ async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch): assert response.status_code == 200 assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads) assert released_by_the_loop and all(released_by_the_loop) + + +@pytest.mark.parametrize( + ("api_base", "route", "expected"), + [ + ("https://gw.test/p/", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p//", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p", "/v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p/", "/v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p/v1/messages", "v1/messages", "https://gw.test/p/v1/messages/v1/messages"), + ("https://api.test/v1/evals/cancel", "cancel", "https://api.test/v1/evals/cancel/cancel"), + ("http://embeddings", "embeddings", "http://embeddings/embeddings"), + ("http:/", "v1/files", "http://v1/files"), + ("https://gw.test/p/?x=1", "v1/messages", "https://gw.test/p?x=1/v1/messages"), + ], + ids=( + "trailing-slash-base", + "repeated-trailing-slashes", + "base-without-slash", + "route-with-leading-slash", + "slashes-on-both-sides", + "base-already-ends-in-route", + "id-equal-to-the-route", + "host-named-like-the-route", + "host-stripped-as-a-suffix-leaves-only-the-scheme", + "slash-before-the-query-matches-the-base-without-it", + ), +) +def test_join_api_base_always_appends_the_route_after_exactly_one_slash( + api_base: str, route: str, expected: str +) -> None: + assert join_api_base(api_base, route) == expected + + +@pytest.mark.parametrize("api_base", ["https://", "http://", "https:///"]) +@pytest.mark.parametrize("route", ["v1/messages", "models/gemini-flash:generateContent"]) +@pytest.mark.parametrize("join", [join_api_base, ensure_api_route], ids=["join_api_base", "ensure_api_route"]) +def test_a_base_that_is_only_a_scheme_never_turns_the_route_into_the_host( + api_base: str, route: str, join: Callable[[str, str], str] +) -> None: + assert urlsplit(join(api_base, route)).netloc == "" + + +@pytest.mark.parametrize( + ("api_base", "route", "expected"), + [ + ("https://gw.test/p/v1/messages", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p/v1/messages/", "/v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/x/listen", "listen", "https://gw.test/x/listen"), + ("https://gw.test/p/", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p", "/v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/research", "search", "https://gw.test/research/search"), + ("http://embeddings", "embeddings", "http://embeddings/embeddings"), + ("http://embeddings:8080/", "embeddings", "http://embeddings:8080/embeddings"), + ("http:/", "anthropic/v1/messages", "http://anthropic/v1/messages"), + ("https://gw.test/p/v1/messages/?beta=true", "v1/messages", "https://gw.test/p/v1/messages?beta=true"), + ("https://gw.test/p/v1/messages?beta=true", "v1/messages", "https://gw.test/p/v1/messages?beta=true"), + ], + ids=( + "path-already-ends-in-route", + "path-ends-in-route-before-trailing-slash", + "last-segment-is-the-route", + "trailing-slash-base", + "route-with-leading-slash", + "path-ends-in-lookalike-suffix", + "host-named-like-the-route", + "host-with-port-named-like-the-route", + "host-stripped-as-a-suffix-leaves-only-the-scheme", + "route-path-with-a-slash-before-the-query", + "route-path-before-the-query", + ), +) +def test_ensure_api_route_appends_the_route_only_when_the_url_path_lacks_it( + api_base: str, route: str, expected: str +) -> None: + assert ensure_api_route(api_base, route) == expected diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py index 785e1f1096b..6951253349e 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py @@ -19,12 +19,14 @@ import json import os import sys import threading +from collections.abc import Callable, Mapping, Sequence from datetime import datetime -from collections.abc import Callable -from typing import List, Optional +from typing import Final, Iterator, List, Optional from unittest.mock import MagicMock +import httpx import pytest +from pydantic import TypeAdapter import litellm @@ -33,12 +35,13 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, _mid_stream_error_sse_event, ) from litellm.llms.bedrock.common_utils import BedrockError -from litellm.types.utils import Delta, StreamingChoices +from litellm.types.utils import Delta, ModelResponseStream, StandardLoggingPayload, StreamingChoices, Usage def _make_chunk(delta: Delta, finish_reason: Optional[str] = None) -> MagicMock: @@ -55,14 +58,14 @@ class _AsyncStreamThenRaise: """Yields the given chunks, then raises ``exc`` (mimics a provider stream that terminates mid-response).""" - def __init__(self, items: List[MagicMock], exc: BaseException): + def __init__(self, items: Sequence[object], exc: BaseException) -> None: self._it = iter(items) self._exc = exc - def __aiter__(self): + def __aiter__(self) -> "_AsyncStreamThenRaise": return self - async def __anext__(self): + async def __anext__(self) -> object: try: return next(self._it) except StopIteration: @@ -182,9 +185,73 @@ class _AsyncFailureRecorder(CustomLogger): def __init__(self): super().__init__() self.exceptions: list[BaseException] = [] + self.usages: tuple[tuple[int, int], ...] = () - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: self.exceptions.append(kwargs["exception"]) + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + self.usages = (*self.usages, (payload["prompt_tokens"], payload["completion_tokens"])) + + +class _BoundaryAsyncFailureRecorder(_AsyncFailureRecorder): + def __init__(self) -> None: + super().__init__() + self.started = 0 + self.completed = 0 + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + self.started += 1 + await asyncio.sleep(0.05) + await super().async_log_failure_event(kwargs, response_obj, start_time, end_time) + self.completed += 1 + + +def test_anthropic_adapter_does_not_repeat_custom_stream_failure_dispatch() -> None: + async_recorder: Final = _BoundaryAsyncFailureRecorder() + sync_recorder: Final = _SyncFailureRecorder() + logging_obj: Final = _make_logging_obj("wrapped-custom-stream", async_recorder, sync_recorder) + chunk: Final = ModelResponseStream( + id="chatcmpl-dropped", + object="chat.completion.chunk", + created=1, + model="gpt-5.6", + choices=[StreamingChoices(index=0, delta=Delta(content="partial"), finish_reason=None)], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + inner_stream: Final = CustomStreamWrapper( + completion_stream=_AsyncStreamThenRaise([chunk], httpx.ReadError("stream dropped")), + model="gpt-5.6", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper: Final = AnthropicStreamWrapper( + completion_stream=inner_stream, + model="openai/gpt-5.6", + litellm_logging_obj=logging_obj, + ) + + events: Final = asyncio.run(_drain_sse(wrapper)) + + assert async_recorder.started == 1 + assert async_recorder.completed == 1 + assert sync_recorder.called.wait(timeout=5), "sync failure callback never ran" + assert _parse_sse(events[-1])[0] == "error" + assert len(async_recorder.exceptions) == 1 + assert len(sync_recorder.exceptions) == 1 + assert async_recorder.usages == ((30, 1),) + assert sync_recorder.exceptions[0] == async_recorder.exceptions[0] class _SyncFailureRecorder: diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_sse_stream.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_sse_stream.py index fbbbc579d94..ec57636a17b 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_sse_stream.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_sse_stream.py @@ -8,8 +8,11 @@ wrapper's chunks/messages/model so disconnect-time partial billing can read them from typing import Final from unittest.mock import MagicMock +import httpx import pytest +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicSSEStream, AnthropicStreamWrapper, @@ -85,3 +88,67 @@ def test_sse_stream_exposes_chunks_messages_and_model(): assert stream.messages == [{"role": "user", "content": "hi"}] chunks: Final = stream.chunks assert isinstance(chunks, list) and len(chunks) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("started", [False, True]) +async def test_sse_stream_close_releases_http_response_before_upstream_body(started: bool) -> None: + response: Final = httpx.Response(200, stream=httpx.ByteStream(b"unread upstream body")) + upstream: Final = CustomStreamWrapper( + completion_stream=response, + model="test-model", + custom_llm_provider="openai", + logging_obj=Logging( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=1.0, + litellm_call_id="disconnect-test", + function_id="test", + ), + ) + stream: Final = AnthropicSSEStream(AnthropicStreamWrapper(completion_stream=upstream, model="test-model")) + if started: + assert (await anext(stream)).startswith(b"event: message_start") + assert not response.is_stream_consumed + assert not response.is_closed + + await stream.aclose() + await stream.aclose() + + assert response.is_closed + assert not response.is_stream_consumed + with pytest.raises(StopAsyncIteration): + await anext(stream) + + +class _CloseableAsyncStream(_AsyncStream): + def __init__(self, items: list[MagicMock]) -> None: + super().__init__(items) + self.close_calls = 0 + + async def aclose(self) -> None: + self.close_calls += 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("finish_stream", [False, True]) +async def test_sse_stream_releases_upstream_once_after_content(finish_stream: bool) -> None: + upstream: Final = _CloseableAsyncStream( + [_make_chunk(Delta(content="hello")), _make_chunk(Delta(), finish_reason="stop")] + ) + stream: Final = AnthropicSSEStream(AnthropicStreamWrapper(completion_stream=upstream, model="test-model")) + events: Final[list[bytes]] = [] + async for event in stream: + events.append(event) + if not finish_stream and b"event: content_block_delta" in event: + break + assert any(b'"text":"hello"' in event.replace(b" ", b"") for event in events) + if finish_stream: + assert events[-1].startswith(b"event: message_stop") + + await stream.aclose() + await stream.aclose() + + assert upstream.close_calls == 1 diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 1f02b3aaa5a..f4920505722 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -2033,6 +2033,7 @@ async def test_anthropic_messages_with_extra_headers(): ], "model": "claude-haiku-4-5-20251001", "stop_reason": "end_turn", + "stop_sequence": None, "usage": {"input_tokens": 10, "output_tokens": 20}, } @@ -2096,6 +2097,7 @@ async def test_anthropic_messages_with_thinking(): ], "model": "claude-haiku-4-5-20251001", "stop_reason": "end_turn", + "stop_sequence": None, "usage": {"input_tokens": 10, "output_tokens": 20}, } @@ -2210,6 +2212,7 @@ async def test_anthropic_messages_bedrock_dynamic_region(): "content": [{"type": "text", "text": "This is a mock response"}], "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "stop_reason": "end_turn", + "stop_sequence": None, "usage": {"input_tokens": 10, "output_tokens": 20}, } diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py index 7fbec3ba12e..457d5d5e1c8 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py @@ -1,17 +1,23 @@ from collections.abc import Mapping +from datetime import datetime, timezone from pathlib import Path from typing import Final from unittest.mock import patch +import httpx import pytest +from pydantic import JsonValue import litellm +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.anthropic.wif import get_anthropic_wif_token from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig +from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig from litellm.llms.tencent.messages.transformation import TencentAnthropicMessagesConfig +from litellm.types.llms.anthropic import AnthropicFinishReason from litellm.types.router import GenericLiteLLMParams from tests.unit.llms.anthropic.test_anthropic_wif import ( ScriptedPoster, @@ -27,6 +33,52 @@ _WIF_PARAMS: Final[dict] = { } +@pytest.mark.parametrize("config", (AnthropicMessagesConfig(), OpenAILikeAnthropicMessagesConfig()), ids=type) +@pytest.mark.parametrize( + ("stop_reason", "stop_fields", "expected_stop_sequence"), + ( + pytest.param("end_turn", {}, None, id="end_turn_missing"), + pytest.param("max_tokens", {}, None, id="max_tokens_missing"), + pytest.param("tool_use", {}, None, id="tool_use_missing"), + pytest.param("end_turn", {"stop_sequence": None}, None, id="explicit_null"), + pytest.param("stop_sequence", {"stop_sequence": ", 7"}, ", 7", id="matched_stop"), + ), +) +def test_native_response_defaults_stop_sequence_without_changing_other_fields( + config: AnthropicMessagesConfig, + stop_reason: AnthropicFinishReason, + stop_fields: Mapping[str, str | None], + expected_stop_sequence: str | None, +) -> None: + upstream: Final[dict[str, JsonValue]] = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "test-model", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": stop_reason, + "usage": {"input_tokens": 19, "output_tokens": 42}, + **stop_fields, + } + logging_obj: Final = Logging( + model="test-model", + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="test-call", + function_id="test-function", + ) + response: Final = config.transform_anthropic_messages_response( + model="test-model", raw_response=httpx.Response(200, json=upstream), logging_obj=logging_obj + ) + + assert response == { + **upstream, + "stop_sequence": expected_stop_sequence, + } + + def test_workload_identity_allowed_for_anthropic() -> None: assert AnthropicMessagesConfig()._allows_workload_identity is True 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 fd9aa75be74..4ed9054f3d3 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 @@ -26,7 +26,9 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( is_provider_error_chunk, parse_anthropic_error_event, ) +from litellm.llms.anthropic.pass_through.messages.transformation import AnthropicMessagesConfig from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator +from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig from litellm.types.utils import Choices, Message, ModelResponse @@ -1733,6 +1735,49 @@ async def _drained(stream: AsyncIterator[bytes]) -> list[bytes]: return [chunk async for chunk in stream] +@pytest.mark.asyncio +@pytest.mark.parametrize("config", (AnthropicMessagesConfig(), OpenAILikeAnthropicMessagesConfig()), ids=type) +@pytest.mark.parametrize("stop_sequence", (None, ", 7"), ids=("null", "matched_stop")) +async def test_native_stream_preserves_nullable_stop_fields_across_chunk_boundaries( + config: AnthropicMessagesConfig, stop_sequence: str | None +) -> None: + logging_obj: Final = _make_logging_obj("native_stream_stop_fields") + events: Final = ( + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": logging_obj.model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 19, "output_tokens": 0}, + }, + }, + { + "type": "message_delta", + "delta": { + "stop_reason": "end_turn" if stop_sequence is None else "stop_sequence", + "stop_sequence": stop_sequence, + }, + "usage": {"output_tokens": 0}, + }, + {"type": "message_stop"}, + ) + upstream: Final = b"".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) + response: Final = httpx.Response(200, content=_chunked(upstream, 7)) + stream: Final = config.get_async_streaming_response_iterator( + model=logging_obj.model, + httpx_response=response, + request_body={"model": logging_obj.model}, + litellm_logging_obj=logging_obj, + ) + + assert b"".join(await _drained(stream)) == upstream + + def _data_events(frames: Sequence[bytes]) -> list[Mapping[str, object]]: return [ json.loads(line[len(b"data:") :]) for frame in frames for line in frame.splitlines() if line.startswith(b"data:") diff --git a/tests/unit/llms/anthropic/test_message_sanitization.py b/tests/unit/llms/anthropic/test_message_sanitization.py index 7afa60baf7e..0afd443d07a 100644 --- a/tests/unit/llms/anthropic/test_message_sanitization.py +++ b/tests/unit/llms/anthropic/test_message_sanitization.py @@ -7,6 +7,8 @@ B. Orphaned tool_result without matching tool_use C. Empty text content """ +import copy +import json import pytest import sys import os @@ -16,6 +18,8 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../. import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( + _EMPTY_TEXT_PLACEHOLDER as PLACEHOLDER, + _sanitize_empty_text_content, sanitize_messages_for_tool_calling, anthropic_messages_pt, ) @@ -24,16 +28,9 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( class TestMessageSanitization: """Test message sanitization for tool calling scenarios""" - def setup_method(self): - """Setup for each test""" - # Save original modify_params value - self.original_modify_params = litellm.modify_params - litellm.modify_params = True - - def teardown_method(self): - """Cleanup after each test""" - # Restore original modify_params value - litellm.modify_params = self.original_modify_params + @pytest.fixture(autouse=True) + def _modify_params_on(self, monkeypatch): + monkeypatch.setattr(litellm, "modify_params", True) def test_case_a_orphaned_tool_call_single(self): """ @@ -255,11 +252,11 @@ class TestMessageSanitization: assert sanitized[3]["content"] == "[System: Empty message content sanitised to satisfy protocol]" assert sanitized[4]["role"] == "assistant" - def test_modify_params_false_no_sanitization(self): + def test_modify_params_false_no_sanitization(self, monkeypatch): """ Test that sanitization is skipped when modify_params=False """ - litellm.modify_params = False + monkeypatch.setattr(litellm, "modify_params", False) messages = [ {"role": "user", "content": ""}, @@ -287,7 +284,6 @@ class TestMessageSanitization: """ Test that sanitization is integrated into anthropic_messages_pt """ - litellm.modify_params = True messages = [ {"role": "user", "content": "What is the weather in Nashik?"}, @@ -316,7 +312,7 @@ class TestMessageSanitization: assert result[0]["role"] == "user" assert result[1]["role"] == "assistant" - def test_empty_string_content_sanitized_without_modify_params(self): + def test_empty_string_content_sanitized_without_modify_params(self, monkeypatch): """ Regression: An empty user message ({"role": "user", "content": ""}) must be rewritten to a non-empty placeholder *before* it reaches Anthropic, @@ -324,7 +320,7 @@ class TestMessageSanitization: "messages: text content blocks must be non-empty" Reproduces a real failure from the pr-review agent (pydantic-ai). """ - litellm.modify_params = False + monkeypatch.setattr(litellm, "modify_params", False) messages = [ {"role": "user", "content": "First message"}, @@ -344,13 +340,13 @@ class TestMessageSanitization: assert block["text"].strip() != "" assert text_blocks[2]["text"] == ("[System: Empty message content sanitised to satisfy protocol]") - def test_empty_text_block_in_list_content_sanitized(self): + def test_empty_text_block_in_list_content_sanitized(self, monkeypatch): """ Same regression for the list-of-blocks form: {"role": "user", "content": [{"type": "text", "text": ""}]} Empty text *blocks* must be rewritten too, regardless of modify_params. """ - litellm.modify_params = False + monkeypatch.setattr(litellm, "modify_params", False) messages = [ { @@ -372,12 +368,12 @@ class TestMessageSanitization: for block in text_blocks[1:]: assert block["text"].strip() != "" - def test_non_empty_content_unchanged_without_modify_params(self): + def test_non_empty_content_unchanged_without_modify_params(self, monkeypatch): """ Sanity check: when nothing is empty, the messages flow through unchanged even with modify_params disabled. """ - litellm.modify_params = False + monkeypatch.setattr(litellm, "modify_params", False) messages = [ {"role": "user", "content": "Hello"}, @@ -393,6 +389,270 @@ class TestMessageSanitization: assert result[1]["content"][0]["text"] == "Hi there" assert result[2]["content"][0]["text"] == "How are you?" + @pytest.mark.parametrize("modify_params", [True, False]) + @pytest.mark.parametrize( + "empty_content", + [None, "", " \n \t ", [], [{"type": "text", "text": ""}]], + ids=["none", "empty_str", "whitespace_str", "empty_list", "empty_text_block"], + ) + def test_tool_call_only_assistant_turn_never_gets_placeholder(self, monkeypatch, modify_params, empty_content): + """ + Regression for the placeholder leaking into tool-call-only assistant turns. + + An OpenAI-format assistant turn that carries tool_calls but no text is the normal + shape emitted by coding agents. Anthropic accepts an assistant turn whose content + is just tool_use, so there is nothing to substitute; injecting the placeholder made + it look like the assistant had said it, and the model echoed it on later turns. + """ + monkeypatch.setattr(litellm, "modify_params", modify_params) + + messages = [ + {"role": "user", "content": "list the files"}, + { + "role": "assistant", + "content": empty_content, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run_bash", "arguments": '{"cmd": "ls"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "a.txt"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assistant_turns = [m for m in result if m["role"] == "assistant"] + assert len(assistant_turns) == 1 + blocks = assistant_turns[0]["content"] + + assert [b["type"] for b in blocks] == ["tool_use"] + assert blocks[0]["id"] == "call_1" + assert blocks[0]["name"] == "run_bash" + + assert PLACEHOLDER not in json.dumps(result) + + def test_empty_text_dropped_alongside_non_text_blocks(self): + """ + An assistant turn already carrying a non-text block has enough content for + Anthropic once the empty text block is dropped, so the placeholder must not be + substituted in its place. + """ + + message = { + "role": "assistant", + "content": [ + {"type": "text", "text": " "}, + {"type": "thinking", "thinking": "weighing options", "signature": "sig"}, + ], + } + + sanitized = _sanitize_empty_text_content(message) + + assert [b["type"] for b in sanitized["content"]] == ["thinking"] + assert PLACEHOLDER not in json.dumps(sanitized) + + @pytest.mark.parametrize( + "sibling_block", + [ + {"type": "thinking", "thinking": "", "signature": "sig"}, + {"type": "redacted_thinking", "data": "opaque"}, + ], + ids=["signed_empty_thinking", "redacted_thinking"], + ) + def test_empty_text_dropped_alongside_kept_thinking_block(self, sibling_block): + """ + The converter keeps signed thinking blocks even with empty text, and keeps + redacted_thinking blocks, so the turn still reaches Anthropic without the + empty text and the placeholder must not be substituted. + """ + + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": " "}, sibling_block]}, + {"role": "user", "content": "and now?"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert [m["role"] for m in result] == ["user", "assistant", "user"] + assert [b["type"] for b in result[1]["content"]] == [sibling_block["type"]] + assert PLACEHOLDER not in json.dumps(result) + + def test_placeholder_kept_when_only_tool_call_is_dropped_tool_search(self, monkeypatch): + """ + A tool search call whose result is not replayed is dropped by the converter, so + it does not count as content: the empty text keeps its placeholder instead of + the assistant turn vanishing from the request. modify_params stays off so no + dummy tool result answers the call and keeps it. + """ + monkeypatch.setattr(litellm, "modify_params", False) + + messages = [ + {"role": "user", "content": "find a tool"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "srvtoolu_1", + "type": "function", + "function": {"name": "tool_search_tool_regex", "arguments": '{"query": "x"}'}, + } + ], + }, + {"role": "user", "content": "go on"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert [m["role"] for m in result] == ["user", "assistant", "user"] + assert result[1]["content"] == [{"type": "text", "text": PLACEHOLDER}] + + @pytest.mark.parametrize( + "sibling_block", + [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}}, + {"type": "thinking", "thinking": ""}, + {"type": "redacted_thinking", "data": " "}, + {"type": "unknown_future_block"}, + {"no_type_key": "malformed"}, + ], + ids=["image_url", "unsignable_thinking", "blank_redacted_thinking", "unknown_type", "missing_type"], + ) + def test_empty_text_kept_when_sibling_block_would_not_survive(self, sibling_block): + """ + Dropping empty text is only safe when something else in the turn actually reaches + Anthropic. Blocks the assistant conversion discards (assistant-side images, + signature-less thinking, unrecognised types) would leave an empty turn that gets + dropped from the request entirely, so the placeholder has to stay. + """ + + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": " "}, sibling_block]}, + {"role": "user", "content": "and now?"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-haiku-4-5", llm_provider="anthropic") + + assert [m["role"] for m in result] == ["user", "assistant", "user"] + assert PLACEHOLDER in json.dumps(result) + + def test_assistant_text_preserved_alongside_tool_calls(self): + """ + Dropping empty text must not drop real text: an assistant turn that says something + *and* calls a tool keeps both blocks. + """ + + messages = [ + {"role": "user", "content": "list the files"}, + { + "role": "assistant", + "content": "Let me check.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run_bash", "arguments": '{"cmd": "ls"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "a.txt"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assistant_turns = [m for m in result if m["role"] == "assistant"] + blocks = assistant_turns[0]["content"] + + assert [b["type"] for b in blocks] == ["text", "tool_use"] + assert blocks[0]["text"] == "Let me check." + assert PLACEHOLDER not in json.dumps(result) + + @pytest.mark.parametrize( + "messages", + [ + [{"role": "user", "content": ""}], + [{"role": "user", "content": [{"type": "text", "text": " "}]}], + [{"role": "user", "content": "hi"}, {"role": "assistant", "content": ""}], + ], + ids=["user_str", "user_block", "assistant_without_tool_calls"], + ) + def test_placeholder_still_used_when_turn_would_be_empty(self, messages): + """ + The placeholder is still required where dropping the empty text would leave the + turn with no content at all, which Anthropic also rejects. + """ + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert PLACEHOLDER in json.dumps(result) + for message in result: + assert message["content"] + + def test_no_empty_text_block_or_content_reaches_anthropic(self): + """ + The invariant Anthropic actually enforces, over a conversation mixing every shape: + no empty content list and no empty text block may survive the conversion. + """ + + messages = [ + {"role": "user", "content": ""}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run_bash", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "out"}, + {"role": "assistant", "content": [{"type": "text", "text": " "}]}, + {"role": "user", "content": [{"type": "text", "text": "thanks"}]}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + for message in result: + content = message["content"] + assert content + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + assert block["text"].strip() + + def test_sanitization_does_not_mutate_caller_messages(self): + """ + Callers reuse their message list across retries and fallbacks, so sanitization has + to stay copy-on-write. + """ + + messages = [ + {"role": "user", "content": "go"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run_bash", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "out"}, + ] + before = copy.deepcopy(messages) + + anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert messages == before + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unit/llms/aws_polly/text_to_speech/test_aws_polly_text_to_speech_transformation.py b/tests/unit/llms/aws_polly/text_to_speech/test_aws_polly_text_to_speech_transformation.py index 3826b22af34..ace82e517ad 100644 --- a/tests/unit/llms/aws_polly/text_to_speech/test_aws_polly_text_to_speech_transformation.py +++ b/tests/unit/llms/aws_polly/text_to_speech/test_aws_polly_text_to_speech_transformation.py @@ -1,6 +1,10 @@ -import litellm +import json +from unittest.mock import MagicMock, patch + +import httpx import pytest -from unittest.mock import MagicMock + +import litellm @pytest.mark.asyncio @@ -134,6 +138,34 @@ async def test_aws_polly_tts_with_native_voice(): assert request_body.get("TextType", "text") == "text" +@pytest.mark.asyncio +@pytest.mark.usefixtures("fake_provider_credentials") +@pytest.mark.parametrize("engine", ["standard", "neural", "long-form", "generative"]) +async def test_aws_polly_tts_sends_engine_from_model(engine: str) -> None: + """ + The engine in the model name must reach Polly. speech() strips the + provider prefix before the config reads it, which used to turn every + engine into the neural default. + """ + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.content = b"fake_audio_data" + mock_httpx_response.status_code = 200 + mock_httpx_response.headers = {"content-type": "audio/mpeg"} + + with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") as mock_post: + mock_post.return_value = mock_httpx_response + + await litellm.aspeech( + model=f"aws_polly/{engine}", + voice="Joanna", + input="Hello", + aws_region_name="us-east-1", + ) + + request_body = json.loads(mock_post.call_args.kwargs["data"]) + assert request_body["Engine"] == engine + + @pytest.mark.asyncio @pytest.mark.usefixtures("fake_provider_credentials") async def test_aws_polly_tts_with_openai_voice_mapping(): diff --git a/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py b/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py index 4778b6296fc..5f195b7a478 100644 --- a/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py +++ b/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py @@ -1,6 +1,8 @@ import json from typing import Final +import pytest + from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToSpeechConfig @@ -17,3 +19,17 @@ def test_installed_botocore_signs_the_speech_request() -> None: authorization: Final = headers["Authorization"] assert "/us-west-2/polly/aws4_request" in authorization assert json.loads(body) == {"Text": "ping", "VoiceId": "Joanna"} + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("standard", "standard"), + ("Long-Form", "long-form"), + ("aws_polly/generative", "generative"), + ("aws_polly", "neural"), + ("unknown", "neural"), + ], +) +def test_extract_engine_from_model(model: str, expected: str) -> None: + assert AWSPollyTextToSpeechConfig()._extract_engine_from_model(model) == expected diff --git a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index 5c1ff1b46a3..f1db76c1e21 100644 --- a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -740,3 +740,15 @@ async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch, no_ambi ) assert mock_async_realtime.call_args.kwargs["azure_ad_token"] == "deployment-entra-token" + + +@pytest.mark.parametrize("protocol", [None, "GA"]) +def test_construct_url_with_a_trailing_slash_api_base_matches_the_url_without_it(protocol): + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + + def url_for(api_base: str) -> str: + return AzureOpenAIRealtime().construct_url( + api_base=api_base, model="gpt-realtime", api_version="2024-10-01-preview", realtime_protocol=protocol + ) + + assert url_for("https://gw.example.com/azure/") == url_for("https://gw.example.com/azure") diff --git a/tests/unit/llms/azure/response/test_azure_transformation.py b/tests/unit/llms/azure/response/test_azure_transformation.py index 7ed98f6c705..d46242fc5f7 100644 --- a/tests/unit/llms/azure/response/test_azure_transformation.py +++ b/tests/unit/llms/azure/response/test_azure_transformation.py @@ -21,6 +21,7 @@ from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +from tests.unit.proxy.conftest import httpx_transport @pytest.mark.serial @@ -100,6 +101,69 @@ def test_get_complete_url(): assert result == expected +@pytest.mark.parametrize( + ("base_path", "expected_path"), + [ + ("/foundry/openai/openai/v1", "/foundry/openai/openai/v1/responses?api-version=v1"), + ("/foundry/openai/openai/v1/", "/foundry/openai/openai/v1/responses?api-version=v1"), + ("/foundry/openai/openai/v1/responses", "/foundry/openai/openai/v1/responses?api-version=v1"), + ("/foundry/openai/openai/v1/responses/", "/foundry/openai/openai/v1/responses?api-version=v1"), + ( + "/openai/v1?redirect=/openai/v1/responses", + "/openai/v1/responses?redirect=%2Fopenai%2Fv1%2Fresponses&api-version=v1", + ), + ( + "/openai/v1?tenant=example&api-version=preview", + "/openai/v1/responses?tenant=example&api-version=preview", + ), + ( + "/openai/v1/responses?api-version=2025-04-01-preview", + "/openai/v1/responses?api-version=2025-04-01-preview", + ), + ( + "/openai/v1/responses?redirect=/openai/responses", + "/openai/v1/responses?redirect=%2Fopenai%2Fresponses&api-version=v1", + ), + ("/foo%2Fbar/openai/v1?tenant=example", "/foo%2Fbar/openai/v1/responses?tenant=example&api-version=v1"), + ( + "/foo%3Fbar/openai/v1/responses?tenant=example", + "/foo%3Fbar/openai/v1/responses?tenant=example&api-version=v1", + ), + ("/openai", "/openai/v1/responses?api-version=v1"), + ("/foundry/openai", "/foundry/openai/v1/responses?api-version=v1"), + ], +) +def test_explicit_v1_api_base_preserves_the_gateway_prefix_and_query(base_path: str, expected_path: str) -> None: + result: Final = AzureOpenAIResponsesAPIConfig().get_complete_url( + api_base=f"https://gateway.example.com{base_path}", litellm_params={"api_version": "v1"} + ) + + assert result == f"https://gateway.example.com{expected_path}" + + +@pytest.mark.parametrize( + ("litellm_params", "expected_version"), + [ + ({}, AZURE_DEFAULT_RESPONSES_API_VERSION), + ({"api_version": "v1"}, "v1"), + ({"api_version": "latest"}, "latest"), + ({"api_version": "preview"}, "preview"), + ({"api_version": "2025-04-01-preview"}, "2025-04-01-preview"), + ({"api_version": "2024-10-21"}, AZURE_DEFAULT_RESPONSES_API_VERSION), + ({"api_version": "2024-10-21", "responses_api_version": "preview"}, "preview"), + ], +) +def test_explicit_v1_api_base_keeps_responses_api_version_selection( + litellm_params: dict[str, str], expected_version: str +) -> None: + result: Final = AzureOpenAIResponsesAPIConfig().get_complete_url( + api_base="https://gateway.example.com/foundry/openai/openai/v1", + litellm_params=litellm_params, + ) + + assert result == f"https://gateway.example.com/foundry/openai/openai/v1/responses?api-version={expected_version}" + + @pytest.mark.serial def test_response_id_path_requests_encode_response_id(): config = AzureOpenAIResponsesAPIConfig() @@ -947,6 +1011,119 @@ def test_dated_ga_api_version_sends_responses_to_the_unset_api_version_url( assert route.calls[0].request.url == _unset_api_version_responses_url() +_COMPLETED_RESPONSE_SSE: Final = "".join( + f"data: {json.dumps(event)}\n\n" + for event in ( + { + "type": "response.output_text.delta", + "sequence_number": 0, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "Done.", + }, + {"type": "response.completed", "sequence_number": 1, "response": _COMPLETED_RESPONSE}, + ) +) + + +@pytest.fixture +def gateway_responses(respx_mock: respx.MockRouter, stream: bool) -> respx.Route: + body: Final = _COMPLETED_RESPONSE_SSE if stream else json.dumps(_COMPLETED_RESPONSE) + content_type: Final = "text/event-stream" if stream else "application/json" + return respx_mock.post( + "https://gateway.example.com/foundry/openai/openai/v1/responses?api-version=v1" + ).mock(return_value=httpx.Response(200, text=body, headers={"content-type": content_type})) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_explicit_v1_base_sends_responses_to_the_gateway(gateway_responses: respx.Route, stream: bool) -> None: + response: Final = litellm.responses( + model="azure/gpt-6.1-sol", + input="hi", + api_base="https://gateway.example.com/foundry/openai/openai/v1", + api_key="azure-test-key", + api_version="v1", + stream=stream, + ) + text: Final = ( + "".join(event.delta for event in response if event.type == "response.output_text.delta") + if stream + else response.output[0].content[0].text + ) + + assert text == "Done." + assert gateway_responses.call_count == 1 + + +@pytest.mark.parametrize("stream", [False, True]) +def test_explicit_v1_base_sends_bridged_chat_to_the_gateway(gateway_responses: respx.Route, stream: bool) -> None: + response: Final = litellm.completion( + model="azure/responses/gpt-6.1-sol", + messages=[{"role": "user", "content": "hi"}], + api_base="https://gateway.example.com/foundry/openai/openai/v1", + api_key="azure-test-key", + api_version="v1", + stream=stream, + ) + text: Final = ( + "".join(chunk.choices[0].delta.content or "" for chunk in response) + if stream + else response.choices[0].message.content + ) + + assert text == "Done." + assert gateway_responses.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.parametrize("stream", [False, True]) +async def test_explicit_v1_base_sends_async_responses_to_the_gateway( + gateway_responses: respx.Route, stream: bool +) -> None: + response: Final = await litellm.aresponses( + model="azure/gpt-6.1-sol", + input="hi", + api_base="https://gateway.example.com/foundry/openai/openai/v1", + api_key="azure-test-key", + api_version="v1", + stream=stream, + ) + text: Final = ( + "".join([event.delta async for event in response if event.type == "response.output_text.delta"]) + if stream + else response.output[0].content[0].text + ) + + assert text == "Done." + assert gateway_responses.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.parametrize("stream", [False, True]) +async def test_explicit_v1_base_sends_async_bridged_chat_to_the_gateway( + gateway_responses: respx.Route, stream: bool +) -> None: + response: Final = await litellm.acompletion( + model="azure/responses/gpt-6.1-sol", + messages=[{"role": "user", "content": "hi"}], + api_base="https://gateway.example.com/foundry/openai/openai/v1", + api_key="azure-test-key", + api_version="v1", + stream=stream, + ) + text: Final = ( + "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + if stream + else response.choices[0].message.content + ) + + assert text == "Done." + assert gateway_responses.call_count == 1 + + _LIFECYCLE_BODIES: Final = { "get": _COMPLETED_RESPONSE, "delete": {"id": "resp_dated", "object": "response", "deleted": True}, diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index b3bd1b8743a..1dc7ff74323 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -1702,6 +1702,24 @@ def test_azure_v1_api_uses_openai_client(api_version): ), f"base_url should contain /openai/v1/, got {async_client.base_url}" +@pytest.mark.parametrize( + "api_base", + [ + "https://test.openai.azure.com", + "https://test.openai.azure.com/", + "https://test.openai.azure.com/openai/v1", + "https://test.openai.azure.com/openai/v1/", + ], +) +@pytest.mark.parametrize("is_async", [False, True]) +def test_azure_v1_api_client_targets_openai_v1_once_whatever_the_api_base_ends_in(api_base: str, is_async: bool): + client: Final = BaseAzureLLM().get_azure_openai_client( + api_key="test-key", api_base=api_base, api_version="v1", _is_async=is_async + ) + + assert str(client.base_url) == "https://test.openai.azure.com/openai/v1/" + + @pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) def test_azure_v1_api_with_azure_ad_token_provider(api_version): """ diff --git a/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py b/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py index 59bec08fca6..3713b1cc895 100644 --- a/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py +++ b/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py @@ -1,4 +1,7 @@ +from unittest.mock import Mock + from litellm.llms.azure.vector_stores.transformation import AzureOpenAIVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig def test_transform_search_vector_store_request_preserves_azure_query_string(): @@ -18,3 +21,29 @@ def test_transform_search_vector_store_request_preserves_azure_query_string(): ) assert url == "https://x.openai.azure.com/openai/vector_stores/vs_1/search?api-version=2024-10-21" + + +def test_get_vector_store_url_puts_id_before_api_version_query(): + config = AzureOpenAIVectorStoreConfig() + api_base = config.get_complete_url( + api_base="https://x.openai.azure.com", + litellm_params={"api_version": "2025-04-01-preview"}, + ) + + url = config.get_vector_store_url(api_base, "vs_abc") + + assert url == "https://x.openai.azure.com/openai/vector_stores/vs_abc?api-version=2025-04-01-preview" + + +def test_get_vector_store_url_without_query_is_plain_path(): + url = AzureOpenAIVectorStoreConfig().get_vector_store_url( + "https://x.openai.azure.com/openai/vector_stores", "vs_abc" + ) + + assert url == "https://x.openai.azure.com/openai/vector_stores/vs_abc" + + +def test_base_get_vector_store_url_appends_id_for_providers_without_query_string(): + url = BaseVectorStoreConfig.get_vector_store_url(Mock(), "https://api.example.test/v1/vector_stores", "vs_abc") + + assert url == "https://api.example.test/v1/vector_stores/vs_abc" diff --git a/tests/unit/llms/azure_ai/agents/test_handler.py b/tests/unit/llms/azure_ai/agents/test_handler.py index f92f5359c49..71466bb172c 100644 --- a/tests/unit/llms/azure_ai/agents/test_handler.py +++ b/tests/unit/llms/azure_ai/agents/test_handler.py @@ -22,10 +22,10 @@ def _httpx_only_transport(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") -def _assistants_api(thread_id: str, answer: str) -> respx.MockRouter: - thread_path: Final = f"/api/projects/litellm-test/threads/{thread_id}" +def _assistants_api(thread_id: str, answer: str, project_path: str = "/api/projects/litellm-test") -> respx.MockRouter: + thread_path: Final = f"{project_path}/threads/{thread_id}" router: Final = respx.mock(assert_all_called=False) - router.post(path="/api/projects/litellm-test/threads").respond( + router.post(path=f"{project_path}/threads").respond( 200, json={"id": thread_id, "object": "thread", "created_at": 1767225600, "metadata": {}} ) router.post(path=f"{thread_path}/messages").respond( @@ -116,6 +116,23 @@ async def test_acompletion_creates_thread_posts_message_runs_and_returns_the_rep assert response._hidden_params["thread_id"] == "thread_new" +@pytest.mark.asyncio +async def test_acompletion_creates_the_thread_under_a_project_named_threads() -> None: + router: Final = _assistants_api("thread_new", "hello", project_path="/api/projects/threads") + + with router: + await litellm.acompletion( + model=f"azure_ai/agents/{AGENT_ID}", + messages=[{"role": "user", "content": "hi"}], + api_base="https://agents-test.services.ai.azure.com/api/projects/threads", + api_key="azure-ad-token", + stream=False, + ) + recorded: Final = _recorded(router) + + assert recorded[0][:2] == ("POST", f"/api/projects/threads/threads?api-version={API_VERSION}") + + @pytest.mark.asyncio async def test_handler_acompletion_with_thread_id_reuses_that_thread() -> None: router: Final = _assistants_api("thread_existing", "Your name is Alice.") 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 5186541a62b..ccac6ffc309 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 @@ -879,6 +879,32 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages(): } +def test_azure_ai_reference_only_tool_result_is_sent_as_empty_text(): + messages: Final = [ + {"role": "user", "content": "load the WebFetch tool"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "ToolSearch", "arguments": "{}"}}], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "tool_reference", "tool_name": "WebFetch"}], + }, + ] + + request: Final = AzureAIStudioConfig().transform_request( + model="fw-glm-5.2", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][2] == {"role": "tool", "tool_call_id": "call_1", "content": ""} + + @pytest.mark.parametrize( "api_base, expected_url", [ diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py index 4b317cff975..8971bd31a60 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py @@ -4,8 +4,9 @@ Tests for Azure AI Anthropic CountTokens transformation. Verifies that the CountTokens API uses the correct authentication headers. """ +from typing import Final - +import pytest from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, @@ -105,3 +106,22 @@ class TestAzureAIAnthropicCountTokensConfig: endpoint == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" ) + + def test_get_count_tokens_endpoint_for_a_resource_named_anthropic(self): + config: Final = AzureAIAnthropicCountTokensConfig() + + endpoint: Final = config.get_count_tokens_endpoint("https://anthropic-res.services.ai.azure.com") + + assert endpoint == "https://anthropic-res.services.ai.azure.com/anthropic/v1/messages/count_tokens" + + @pytest.mark.parametrize( + "api_base", + [ + "https://my-resource.services.ai.azure.com/anthropic/v1/messages", + "https://my-resource.services.ai.azure.com/anthropic/v1/messages/", + ], + ) + def test_get_count_tokens_endpoint_sits_next_to_the_messages_url(self, api_base: str): + endpoint: Final = AzureAIAnthropicCountTokensConfig().get_count_tokens_endpoint(api_base) + + assert endpoint == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" 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 38483b8a1fa..5e0395f84dc 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 @@ -198,6 +198,17 @@ class TestAzureAnthropicMessagesConfig: assert url == "https://test.services.ai.azure.com/anthropic/v1/messages" + def test_get_complete_url_adds_the_anthropic_path_for_a_host_named_anthropic(self): + url: Final = AzureAnthropicMessagesConfig().get_complete_url( + api_base="https://anthropic-res.services.ai.azure.com", + api_key="test-api-key", + model="claude-sonnet-4-5", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://anthropic-res.services.ai.azure.com/anthropic/v1/messages" + def test_get_complete_url_raises_error_when_api_base_missing(self): """Test get_complete_url raises error when api_base is None""" config = AzureAnthropicMessagesConfig() diff --git a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index 053ba477212..413061a732d 100644 --- a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -1269,3 +1269,17 @@ def test_agentcore_runtime_response_becomes_assistant_message( assert json.loads(route.calls[0].request.content)["prompt"] == "Explain machine learning in simple terms" assert response.choices[0].message.content == "Machine learning lets computers learn from data." assert response.choices[0].finish_reason == "stop" + + +def test_get_complete_url_sends_a_non_string_runtime_endpoint_to_that_endpoint(): + vpc_endpoint: Final = "https://vpce-1.bedrock-agentcore.us-west-2.vpce.amazonaws.com" + + url: Final = AmazonAgentCoreConfig().get_complete_url( + api_base=None, + api_key=None, + model="agentcore/arn:aws:bedrock-agentcore:us-west-2:111122223333:runtime/my-agent", + optional_params={"aws_bedrock_runtime_endpoint": httpx.URL(f"{vpc_endpoint}/")}, + litellm_params={}, + ) + + assert url.startswith(f"{vpc_endpoint}/runtimes/") diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py index ffddb09cca7..8cf69ba4cd3 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py @@ -1,5 +1,6 @@ import asyncio import base64 +import contextlib import hashlib import json import logging @@ -18,6 +19,9 @@ import tiktoken import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + HUGGING_FACE_SOURCE, + KIMI_K3_RANKS_FILE, + KIMI_K3_REVISION, KimiK3Encodings, KimiK3TokenizerSource, kimi_k3_anthropic_tokens, @@ -522,3 +526,40 @@ async def test_tokenizer_load_still_running_after_the_wait_is_logged( assert loaded is None assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records) + + +@pytest.mark.parametrize( + ("offline_mode", "disable_download", "asks_hub"), + [("true", None, False), ("false", True, False), ("false", None, True)], +) +def test_offline_mode_and_disabled_downloads_never_ask_hugging_face_for_the_tokenizer( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, offline_mode: str, disable_download: bool | None, asks_hub: bool +): + import huggingface_hub.constants + import huggingface_hub.utils._detect_agent + + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_CACHE", str(tmp_path / "hub")) + monkeypatch.setattr(huggingface_hub.constants, "AGENT_HARNESSES_PATH", str(tmp_path / "agent_harnesses.json")) + monkeypatch.setattr(huggingface_hub.utils._detect_agent, "_registry", None) + monkeypatch.setenv("LITELLM_OFFLINE_MODE", offline_mode) + monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", disable_download) + with respx.mock(assert_all_called=False) as router: + hub: Final = router.route(host="huggingface.co").mock(return_value=httpx.Response(503)) + with contextlib.suppress(Exception): + HUGGING_FACE_SOURCE.ranks_path() + assert hub.called is asks_hub + + +def test_offline_mode_reads_the_tokenizer_from_the_hugging_face_cache(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + import huggingface_hub.constants + + cached: Final = tmp_path / "models--moonshotai--Kimi-K3" / "snapshots" / KIMI_K3_REVISION / KIMI_K3_RANKS_FILE + cached.parent.mkdir(parents=True) + cached.write_bytes(b"cached ranks") + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") + with respx.mock(assert_all_called=False) as router: + hub: Final = router.route(host="huggingface.co").mock(return_value=httpx.Response(503)) + path: Final = HUGGING_FACE_SOURCE.ranks_path() + assert Path(path).read_bytes() == b"cached ranks" + assert not hub.called diff --git a/tests/unit/llms/bedrock/test_cost_calculation.py b/tests/unit/llms/bedrock/test_cost_calculation.py new file mode 100644 index 00000000000..f4718344e7e --- /dev/null +++ b/tests/unit/llms/bedrock/test_cost_calculation.py @@ -0,0 +1,43 @@ +import json +from importlib.resources import files +from typing import Final + +import pytest + +import litellm +from litellm.types.utils import Usage + +CLAUDE_PLATFORM_PREFIX: Final = "bedrock/claude_platform/" +CLAUDE_PLATFORM_KEYS: Final = tuple( + key + for key in json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text()) + if key.startswith(CLAUDE_PLATFORM_PREFIX) +) + + +def test_claude_platform_rows_exist() -> None: + assert CLAUDE_PLATFORM_KEYS + + +@pytest.mark.parametrize("inference_geo", ["global", "us"]) +@pytest.mark.parametrize("key", CLAUDE_PLATFORM_KEYS) +def test_claude_platform_cost_matches_claude_api_cost( + local_model_cost_map: None, key: str, inference_geo: str +) -> None: + usage: Final = Usage( + prompt_tokens=1_000, + completion_tokens=200, + cache_read_input_tokens=300, + cache_creation_input_tokens=250, + inference_geo=inference_geo, + ) + + claude_platform_cost: Final = litellm.cost_per_token( + model=key.removeprefix("bedrock/"), custom_llm_provider="bedrock", usage_object=usage + ) + claude_api_cost: Final = litellm.cost_per_token( + model=key.removeprefix(CLAUDE_PLATFORM_PREFIX), custom_llm_provider="anthropic", usage_object=usage + ) + + assert claude_platform_cost == pytest.approx(claude_api_cost) + assert sum(claude_platform_cost) > 0 diff --git a/tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py b/tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py new file mode 100644 index 00000000000..f9d1dc28d12 --- /dev/null +++ b/tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py @@ -0,0 +1,638 @@ +from unittest.mock import MagicMock + +import pytest + +from litellm.exceptions import AuthenticationError +from litellm.llms.chatgpt.chat.transformation import ChatGPTConfig +from litellm.llms.chatgpt.common_utils import GetAccessTokenError + +GATEWAY_BASE = "https://chatgpt.example/backend-api/codex" + + +def _gateway_authenticator() -> MagicMock: + authenticator = MagicMock() + authenticator.get_api_base.return_value = GATEWAY_BASE + authenticator.get_access_token.return_value = "gateway-access-token" + return authenticator + + +def _provider_info( + config: ChatGPTConfig, api_base: str | None, api_key: str | None +) -> tuple[str | None, str | None, str]: + return config._get_openai_compatible_provider_info( + model="gpt-5.4", api_base=api_base, api_key=api_key, custom_llm_provider="chatgpt" + ) + + +class TestChatGPTChatProviderInfo: + def test_configured_key_uses_configured_base_without_gateway_login(self): + authenticator = _gateway_authenticator() + config = ChatGPTConfig(authenticator=authenticator) + + assert _provider_info(config, "http://127.0.0.1:9999", "configured-key") == ( + "http://127.0.0.1:9999", + "configured-key", + "chatgpt", + ) + authenticator.get_access_token.assert_not_called() + + def test_configured_key_without_base_uses_gateway_base(self): + config = ChatGPTConfig(authenticator=_gateway_authenticator()) + + assert _provider_info(config, None, "configured-key") == (GATEWAY_BASE, "configured-key", "chatgpt") + + def test_gateway_login_is_never_sent_to_a_custom_base(self): + config = ChatGPTConfig(authenticator=_gateway_authenticator()) + + assert _provider_info(config, "http://attacker.example", None) == ( + GATEWAY_BASE, + "gateway-access-token", + "chatgpt", + ) + + def test_missing_gateway_login_raises_authentication_error(self): + authenticator = _gateway_authenticator() + authenticator.get_access_token.side_effect = GetAccessTokenError(status_code=401, message="not logged in") + config = ChatGPTConfig(authenticator=authenticator) + + with pytest.raises(AuthenticationError, match="not logged in"): + _provider_info(config, None, None) + + +import asyncio +import json +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import respx + +import litellm +from litellm.integrations.dotprompt.dotprompt_manager import DotpromptManager +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.chatgpt.chat.transformation import ChatGPTConfig +from litellm.types.router import GenericLiteLLMParams + + +def _write_auth_record(path: Path, token: str, account_id: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "access_token": token, + "account_id": account_id, + "expires_at": 4102444800, + } + ) + ) + + +@pytest.fixture +def auth_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, Path]: + monkeypatch.setattr("time.time", lambda: 1700000000) + default_dir: Final = tmp_path / "default" + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + _write_auth_record(default_dir / "auth.json", "token-default", "acct-default") + file_a: Final = tmp_path / "account-a" / "auth.json" + file_b: Final = tmp_path / "account-b" / "auth.json" + _write_auth_record(file_a, "token-a", "acct-a") + _write_auth_record(file_b, "token-b", "acct-b") + return file_a, file_b + + +class TestChatGPTMultiAccountChat: + def test_validate_environment_uses_per_model_auth_file(self, auth_files: tuple[Path, Path]) -> None: + file_a, file_b = auth_files + config: Final = ChatGPTConfig() + + headers_a: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + messages=[], + optional_params={}, + litellm_params={"chatgpt_auth_file": str(file_a)}, + api_key="token-default", + ) + headers_b: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + messages=[], + optional_params={}, + litellm_params={"chatgpt_auth_file": str(file_b)}, + api_key="token-default", + ) + + assert headers_a["Authorization"] == "Bearer token-a" + assert headers_a["ChatGPT-Account-Id"] == "acct-a" + assert headers_b["Authorization"] == "Bearer token-b" + assert headers_b["ChatGPT-Account-Id"] == "acct-b" + + def test_validate_environment_defaults_to_global_auth_file(self, auth_files: tuple[Path, Path]) -> None: + config: Final = ChatGPTConfig() + + headers: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + messages=[], + optional_params={}, + litellm_params={}, + api_key="token-default", + ) + + assert headers["Authorization"] == "Bearer token-default" + assert headers["ChatGPT-Account-Id"] == "acct-default" + + def test_provider_info_uses_per_model_auth_file(self, auth_files: tuple[Path, Path]) -> None: + file_a, _ = auth_files + config: Final = ChatGPTConfig() + + _, per_model_key, _ = config._get_openai_compatible_provider_info( + "gpt-6.1-sol", None, None, "chatgpt", {"chatgpt_auth_file": str(file_a)} + ) + _, default_key, _ = config._get_openai_compatible_provider_info("gpt-6.1-sol", None, None, "chatgpt") + + assert per_model_key == "token-a" + assert default_key == "token-default" + + def test_get_llm_provider_threads_auth_file(self, auth_files: tuple[Path, Path]) -> None: + file_a, _ = auth_files + + model, provider, dynamic_api_key, _ = get_llm_provider( + model="chatgpt/gpt-6.1-sol", + litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_a)), + ) + + assert model == "gpt-6.1-sol" + assert provider == "chatgpt" + assert dynamic_api_key == "token-a" + + +def _subscription_response(request: httpx.Request) -> httpx.Response: + account: Final = request.headers["ChatGPT-Account-Id"] + assert request.headers["Authorization"] in ( + f"Bearer token-{account.removeprefix('acct-')}", + f"Bearer {_client_token(account)}", + ) + assert "chatgpt_auth_file" not in request.content.decode() + message: Final = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": account, "annotations": []}], + } + response: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [message], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + events: Final = ( + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + { + "type": "response.output_text.delta", + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + "delta": account, + }, + {"type": "response.output_item.done", "output_index": 0, "item": message}, + {"type": "response.completed", "response": response}, + ) + return httpx.Response( + 200, + text="".join(f"data: {json.dumps(event)}\n\n" for event in events), + headers={"content-type": "text/event-stream"}, + ) + + +@pytest.mark.parametrize("blocked_default", [False, True]) +@pytest.mark.parametrize("with_prompt", [False, True]) +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("through_router", [False, True]) +@pytest.mark.asyncio +async def test_requests_use_selected_subscription( + auth_files: tuple[Path, Path], + endpoint: str, + through_router: bool, + with_prompt: bool, + blocked_default: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + default_dir: Final = auth_files[0].parent / "missing-default" + if blocked_default: + default_dir.write_text("ordinary file") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir / "unused")) + if with_prompt: + monkeypatch.setattr( + litellm, + "callbacks", + [ + DotpromptManager( + prompt_id="account-check", + prompt_data={"content": "Saved account prompt", "metadata": {"model": "chatgpt/gpt-6.1-sol"}}, + ) + ], + ) + router: Final = ( + litellm.Router( + model_list=[ + { + "model_name": f"account-{index}", + "litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(path)}, + } + for index, path in enumerate(auth_files) + ] + ) + if through_router + else None + ) + + async def request(index: int, path: Path) -> str: + params: Final = MappingProxyType( + { + "model": f"account-{index}" if router else "chatgpt/gpt-6.1-sol", + **({"chatgpt_auth_file": str(path)} if router is None else {}), + **({"prompt_id": "account-check"} if with_prompt else {}), + } + ) + if endpoint == "chat": + result: Final = await (router.acompletion if router else litellm.acompletion)( + **params, messages=[{"role": "user", "content": "Hello"}] + ) + return result.choices[0].message.content + response: Final = await (router.aresponses if router else litellm.aresponses)( + **params, input=[{"role": "user", "content": "Hello"}] + ) + return response.output[0].content[0].text + + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=_subscription_response) + results: Final = await asyncio.gather(*(request(index, path) for index, path in enumerate(auth_files))) + assert results == ["acct-a", "acct-b"] + if with_prompt: + assert all("Saved account prompt" in call.request.content.decode() for call in transport.calls) + + +def _client_token(account: str) -> str: + import base64 + + payload: Final = ( + base64.urlsafe_b64encode(json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account}}).encode()) + .rstrip(b"=") + .decode() + ) + return f"header.{payload}.signature" + + +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("configured_login", [False, True]) +@pytest.mark.parametrize("with_prompt", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_client_login_keeps_its_account_across_gateway_deployments( + auth_files: tuple[Path, Path], + endpoint: str, + configured_login: bool, + with_prompt: bool, + stream: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(auth_files[0].parent / "no-gateway-login")) + if with_prompt: + monkeypatch.setattr( + litellm, + "callbacks", + [ + DotpromptManager( + prompt_id="client-account-check", + prompt_data={"content": "Saved client prompt", "metadata": {"model": "chatgpt/gpt-6.1-sol"}}, + ) + ], + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "subscriptions", + "litellm_params": { + "model": "chatgpt/gpt-6.1-sol", + **({"chatgpt_auth_file": str(path)} if configured_login else {}), + }, + } + for path in auth_files + ], + num_retries=0, + ) + + def reply(request: httpx.Request) -> httpx.Response: + account: Final = request.headers["ChatGPT-Account-Id"] + assert request.headers["Authorization"] == f"Bearer {_client_token(account)}" + return _subscription_response( + httpx.Request( + "POST", + request.url, + content=request.content, + headers={ + **{key: value for key, value in request.headers.items() if key.lower() != "authorization"}, + "Authorization": f"Bearer token-{account.removeprefix('acct-')}", + }, + ) + ) + + async def call(account: str) -> str: + params: Final = MappingProxyType( + { + "model": "subscriptions", + "stream": stream, + **({"prompt_id": "client-account-check"} if with_prompt else {}), + "provider_specific_header": { + "custom_llm_provider": "chatgpt", + "extra_headers": {"Authorization": f"Bearer {_client_token(account)}"}, + }, + } + ) + response: Final = await ( + router.acompletion(**params, messages=[{"role": "user", "content": "Hello"}]) + if endpoint == "chat" + else router.aresponses(**params, input=[{"role": "user", "content": "Hello"}]) + ) + if stream: + chunks: Final = [part async for part in response] + return ( + "".join(part.choices[0].delta.content or "" for part in chunks) + if endpoint == "chat" + else "".join( + getattr(part, "delta", "") + for part in chunks + if getattr(part, "type", "") == "response.output_text.delta" + ) + ) + return response.choices[0].message.content if endpoint == "chat" else response.output[0].content[0].text + + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=reply) + results: Final = await asyncio.gather(call("acct-client-a"), call("acct-client-b")) + assert results == ["acct-client-a", "acct-client-b"] + assert len(transport.calls) == 2 + + +@pytest.mark.asyncio +async def test_one_client_limit_does_not_cool_down_other_subscriptions( + auth_files: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + {"model_name": "pool", "litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(path)}} + for path in auth_files + ], + num_retries=0, + allowed_fails=0, + cooldown_time=60, + ) + headers: Final = { + "custom_llm_provider": "chatgpt", + "extra_headers": {"Authorization": f"Bearer {_client_token('acct-limited')}"}, + } + with respx.mock as transport: + route: Final = transport.post("https://chatgpt.com/backend-api/codex/responses") + route.respond( + 429, + headers={"retry-after": "60"}, + json={"error": {"message": "Usage limit reached", "type": "rate_limit_error"}}, + ) + for _ in range(2): + with pytest.raises(litellm.RateLimitError): + await router.aresponses(model="pool", input="Hello", provider_specific_header=headers) + await asyncio.sleep(0) + route.mock(side_effect=_subscription_response) + response: Final = await router.aresponses(model="pool", input="Hello") + assert response.output[0].content[0].text in ("acct-a", "acct-b") + + +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("async_call", [False, True]) +@pytest.mark.parametrize("direct_headers", [False, True]) +@pytest.mark.parametrize("configured_key", [False, True]) +@pytest.mark.parametrize( + "target", ["https://example.com/v1", "http://chatgpt.com/backend-api/codex", "https://chatgpt.com.example.com/v1"] +) +@pytest.mark.asyncio +async def test_client_login_cannot_be_redirected_to_another_host( + endpoint: str, target: str, async_call: bool, direct_headers: bool, configured_key: bool +) -> None: + headers: Final = {"Authorization": f"Bearer {_client_token('acct-client')}"} + credential: Final = ( + {"extra_headers": headers} + if direct_headers + else {"provider_specific_header": {"custom_llm_provider": "chatgpt", "extra_headers": headers}} + ) + params: Final = MappingProxyType( + { + "model": "chatgpt/gpt-6.1-sol", + "api_base": target, + **({"api_key": "configured-test-key"} if configured_key else {}), + **credential, + } + ) + + async def request() -> None: + if endpoint == "chat": + if async_call: + await litellm.acompletion(**params, messages=[{"role": "user", "content": "Hello"}]) + else: + litellm.completion(**params, messages=[{"role": "user", "content": "Hello"}]) + elif async_call: + await litellm.aresponses(**params, input="Hello") + else: + litellm.responses(**params, input="Hello") + + with respx.mock as transport: + with pytest.raises(litellm.BadRequestError, match="official HTTPS ChatGPT endpoint"): + await request() + assert len(transport.calls) == 0 + + +@pytest.mark.asyncio +async def test_cross_provider_fallback_does_not_receive_the_chatgpt_login( + auth_files: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "subscription", + "litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(auth_files[0])}, + }, + { + "model_name": "api-fallback", + "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": "sk-example-fallback"}, + }, + ], + fallbacks=[{"subscription": ["api-fallback"]}], + num_retries=0, + ) + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").respond( + 429, json={"error": {"message": "Usage limit reached", "type": "rate_limit_error"}} + ) + fallback: Final = transport.post("https://api.openai.com/v1/responses").respond( + 200, + json={ + "id": "resp_fallback", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [ + { + "id": "msg_fallback", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "fallback", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + response: Final = await router.aresponses( + model="subscription", + input="Hello", + provider_specific_header={ + "custom_llm_provider": "chatgpt", + "extra_headers": { + "Authorization": f"Bearer {_client_token('acct-client')}", + "ChatGPT-Account-Id": "acct-client", + }, + }, + ) + assert response.output[0].content[0].text == "fallback" + assert fallback.calls.last.request.headers["Authorization"] == "Bearer sk-example-fallback" + assert "ChatGPT-Account-Id" not in fallback.calls.last.request.headers + + +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("async_call", [False, True]) +@pytest.mark.parametrize("with_prompt", [False, True]) +@pytest.mark.asyncio +async def test_sdk_extra_headers_use_client_login_without_gateway_auth( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, endpoint: str, async_call: bool, with_prompt: bool +) -> None: + blocked: Final = tmp_path / "blocked" + blocked.write_text("ordinary file") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(blocked / "unused")) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + headers: Final = {"Authorization": f"Bearer {_client_token('acct-client-a')}"} + prompt_params: Final = {"prompt_id": "sdk-client"} if with_prompt else {} + if with_prompt: + monkeypatch.setattr( + litellm, + "callbacks", + [ + DotpromptManager( + prompt_id="sdk-client", + prompt_data={"content": "Hello", "metadata": {"model": "chatgpt/gpt-6-luna"}}, + ) + ], + ) + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=_subscription_response) + if endpoint == "chat": + params: Final = { + "model": "chatgpt/gpt-6.1-sol", + "messages": [{"role": "user", "content": "Hello"}], + "extra_headers": headers, + **prompt_params, + } + result: Final = await litellm.acompletion(**params) if async_call else litellm.completion(**params) + assert result.choices[0].message.content == "acct-client-a" + else: + response: Final = ( + await litellm.aresponses( + model="chatgpt/gpt-6.1-sol", input="Hello", extra_headers=headers, **prompt_params + ) + if async_call + else litellm.responses( + model="chatgpt/gpt-6.1-sol", input="Hello", extra_headers=headers, **prompt_params + ) + ) + assert response.output[0].content[0].text == "acct-client-a" + assert transport.calls.last.request.headers["Authorization"] == headers["Authorization"] + assert len(transport.calls) == 1 + + +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize("credential_source", ["default", "file", "api_key"]) +@pytest.mark.parametrize("authorization_header", [None, "Authorization", "authorization"]) +@pytest.mark.asyncio +async def test_requests_preserve_gateway_token_and_selected_workspace( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + endpoint: str, + sync: bool, + credential_source: str, + authorization_header: str | None, +) -> None: + monkeypatch.setattr("time.time", lambda: 1700000000) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("CHATGPT_AUTH_FILE", str(tmp_path / "default.json")) + token: Final = _client_token("acct-token-default") + auth_file: Final = tmp_path / ("selected.json" if credential_source == "file" else "default.json") + _write_auth_record(auth_file, token, "acct-saved-workspace") + configured_token: Final = _client_token("acct-configured") + expected_token: Final = configured_token if credential_source == "api_key" else token + expected_account: Final = "acct-configured" if credential_source == "api_key" else "acct-saved-workspace" + + def reply(request: httpx.Request) -> httpx.Response: + assert request.headers.get_list("authorization") == [f"Bearer {expected_token}"] + assert request.headers.get_list("chatgpt-account-id") == [expected_account] + assert request.headers["originator"] == "codex_cli_rs" + return _subscription_response( + httpx.Request( + "POST", + request.url, + content=request.content, + headers={ + "Authorization": f"Bearer token-{expected_account.removeprefix('acct-')}", + "ChatGPT-Account-Id": expected_account, + }, + ) + ) + + params: Final = MappingProxyType( + { + "model": "chatgpt/gpt-6.1-sol", + **({"chatgpt_auth_file": str(auth_file)} if credential_source == "file" else {}), + **({"api_key": configured_token} if credential_source == "api_key" else {}), + "extra_headers": { + "originator": "codex_cli_rs", + **({authorization_header: "Bearer sk-unrelated-example"} if authorization_header else {}), + }, + } + ) + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=reply) + if endpoint == "chat": + result: Final = ( + await asyncio.to_thread(litellm.completion, **params, messages=[{"role": "user", "content": "Hello"}]) + if sync + else await litellm.acompletion(**params, messages=[{"role": "user", "content": "Hello"}]) + ) + assert result.choices[0].message.content == expected_account + else: + response: Final = ( + await asyncio.to_thread(litellm.responses, **params, input=[{"role": "user", "content": "Hello"}]) + if sync + else await litellm.aresponses(**params, input=[{"role": "user", "content": "Hello"}]) + ) + assert response.output[0].content[0].text == expected_account diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index e47d7fc702c..d88bdfc0b41 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -4,8 +4,10 @@ Tests for ChatGPT subscription Responses API transformation Source: litellm/llms/chatgpt/responses/transformation.py """ +import base64 import json from collections.abc import Generator +from pathlib import Path from typing import Final from unittest.mock import MagicMock, patch @@ -13,6 +15,7 @@ import httpx import pytest import litellm +from litellm.llms.chatgpt.authenticator import Authenticator from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig from litellm.llms.openai.common_utils import OpenAIError from litellm.main import responses_api_bridge_check @@ -395,3 +398,165 @@ class TestChatGPTResponsesAPITransformation: assert exc_info.value.message == "ChatGPT upstream failed" assert exc_info.value.status_code == 502 + + +def _fake_chatgpt_oauth_token(account_id: str) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode( + json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}}).encode() + ) + .rstrip(b"=") + .decode() + ) + return f"{header}.{payload}.fake-signature" + + +def _gateway_authenticator() -> MagicMock: + authenticator = MagicMock() + authenticator.get_access_token.return_value = "gateway-access-token" + authenticator.get_account_id.return_value = "acct-gateway" + return authenticator + + +def _config_with(authenticator: MagicMock) -> ChatGPTResponsesAPIConfig: + return ChatGPTResponsesAPIConfig(authenticator=authenticator) + + +class TestChatGPTClientCredentialForwarding: + def test_client_chatgpt_oauth_wins_over_configured_key_and_gateway_login(self): + authenticator = _gateway_authenticator() + token = _fake_chatgpt_oauth_token("acct-alice") + + headers = _config_with(authenticator).validate_environment( + headers={"Authorization": f"Bearer {token}", "originator": "codex_cli_rs"}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key", litellm_session_id="session-1"), + ) + + assert headers["Authorization"] == f"Bearer {token}" + assert headers["ChatGPT-Account-Id"] == "acct-alice" + assert headers["originator"] == "codex_cli_rs" + assert headers["accept"] == "text/event-stream" + assert [name for name in headers if name.lower() == "authorization"] == ["Authorization"] + authenticator.get_access_token.assert_not_called() + + def test_explicit_account_header_wins_over_jwt_claim(self): + token = _fake_chatgpt_oauth_token("acct-from-jwt") + + headers = _config_with(_gateway_authenticator()).validate_environment( + headers={"authorization": f"Bearer {token}", "chatgpt-account-id": "acct-from-header"}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers["ChatGPT-Account-Id"] == "acct-from-header" + assert [name for name in headers if name.lower() == "chatgpt-account-id"] == ["ChatGPT-Account-Id"] + + @pytest.mark.parametrize( + "bearer", ["Bearer sk-proj-fake-platform-key", "Bearer not-a-jwt", "Bearer a.bm90LWpzb24.c"] + ) + def test_non_chatgpt_bearer_falls_back_to_configured_key(self, bearer, tmp_path): + authenticator = Authenticator(auth_file=str(tmp_path / "missing-default.json")) + + headers = ChatGPTResponsesAPIConfig(authenticator=authenticator).validate_environment( + headers={"Authorization": bearer}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key"), + ) + + assert "ChatGPT-Account-Id" not in headers + assert headers["Authorization"] == "Bearer configured-gateway-key" + + def test_configured_key_is_used_without_client_credential(self): + authenticator = _gateway_authenticator() + + headers = _config_with(authenticator).validate_environment( + headers={}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key"), + ) + + assert headers["Authorization"] == "Bearer configured-gateway-key" + authenticator.get_access_token.assert_not_called() + + def test_gateway_login_is_used_without_client_credential_or_configured_key(self): + headers = _config_with(_gateway_authenticator()).validate_environment( + headers={}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers["Authorization"] == "Bearer gateway-access-token" + assert headers["ChatGPT-Account-Id"] == "acct-gateway" + + def test_two_clients_never_share_a_credential(self): + config = _config_with(_gateway_authenticator()) + alice = _fake_chatgpt_oauth_token("acct-alice") + bob = _fake_chatgpt_oauth_token("acct-bob") + + forwarded = [ + config.validate_environment( + headers={"Authorization": f"Bearer {token}"}, model="gpt-5.4", litellm_params=GenericLiteLLMParams() + ) + for token in (alice, bob, alice) + ] + + assert [h["Authorization"] for h in forwarded] == [f"Bearer {alice}", f"Bearer {bob}", f"Bearer {alice}"] + assert [h["ChatGPT-Account-Id"] for h in forwarded] == ["acct-alice", "acct-bob", "acct-alice"] + + +class TestChatGPTResponsesMultiAccount: + @pytest.fixture(autouse=True) + def fixed_clock(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("time.time", lambda: 1700000000) + + @staticmethod + def _write_auth_record(path: Path, token: str, account_id: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "access_token": token, + "account_id": account_id, + "expires_at": 4102444800, + } + ) + ) + + def test_validate_environment_uses_per_model_auth_file( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + default_dir: Final = tmp_path / "default" + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + self._write_auth_record(default_dir / "auth.json", "token-default", "acct-default") + file_a: Final = tmp_path / "account-a" / "auth.json" + file_b: Final = tmp_path / "account-b" / "auth.json" + self._write_auth_record(file_a, "token-a", "acct-a") + self._write_auth_record(file_b, "token-b", "acct-b") + + config: Final = ChatGPTResponsesAPIConfig() + + headers_a: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_a)), + ) + headers_b: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_b)), + ) + headers_default: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers_a["Authorization"] == "Bearer token-a" + assert headers_a["ChatGPT-Account-Id"] == "acct-a" + assert headers_b["Authorization"] == "Bearer token-b" + assert headers_b["ChatGPT-Account-Id"] == "acct-b" + assert headers_default["Authorization"] == "Bearer token-default" + assert headers_default["ChatGPT-Account-Id"] == "acct-default" diff --git a/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py index 43a4c26b6d2..2e9e3462bee 100644 --- a/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py +++ b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py @@ -2,17 +2,24 @@ import base64 import json import time from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final from unittest.mock import MagicMock, mock_open, patch import pytest +import httpx +import respx import litellm from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS, HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS from litellm.llms.chatgpt.authenticator import ( TOKEN_REFRESH_TIMEOUT_SECONDS, Authenticator, + get_cached_authenticator, + get_chatgpt_auth_file, ) from litellm.llms.chatgpt.common_utils import GetAccessTokenError +from litellm.types.router import GenericLiteLLMParams def _make_jwt(payload: dict) -> str: @@ -161,3 +168,144 @@ class TestChatGPTAuthenticator: assert account_id == "acct-123" mock_write.assert_called_once() assert mock_write.call_args[0][0]["account_id"] == "acct-123" + + def test_get_account_id_prefers_stored_value_without_rewriting_file(self, authenticator): + id_token = _make_jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-from-token"}}) + auth_data = json.dumps({"account_id": "acct-stored", "id_token": id_token}) + + with ( + patch("builtins.open", mock_open(read_data=auth_data)), + patch.object(authenticator, "_write_auth_file") as mock_write, + ): + assert authenticator.get_account_id() == "acct-stored" + mock_write.assert_not_called() + + def test_get_account_id_without_login_returns_none(self, authenticator): + with ( + patch("builtins.open", mock_open(read_data="{}")), + patch.object(authenticator, "_write_auth_file") as mock_write, + ): + assert authenticator.get_account_id() is None + mock_write.assert_not_called() + + +def _write_auth_record(path: Path, token: str, account_id: str) -> None: + path.write_text( + json.dumps( + { + "access_token": token, + "account_id": account_id, + "expires_at": 4102444800, + } + ) + ) + + +class TestChatGPTMultiAccountAuthenticator: + @pytest.fixture(autouse=True) + def fixed_clock(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("time.time", lambda: 1700000000) + + def test_explicit_auth_file_overrides_env(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "env-dir")) + monkeypatch.setenv("CHATGPT_AUTH_FILE", "env.json") + custom_file: Final = tmp_path / "account-a" / "auth.json" + + authenticator: Final = Authenticator(auth_file=str(custom_file)) + + assert authenticator.auth_file == str(custom_file) + assert authenticator.token_dir == str(custom_file.parent) + assert not custom_file.parent.exists() + authenticator._write_auth_file({"access_token": "token-a"}) + assert json.loads(custom_file.read_text())["access_token"] == "token-a" + + def test_default_auth_file_from_env(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "env-dir")) + monkeypatch.setenv("CHATGPT_AUTH_FILE", "env.json") + + authenticator: Final = Authenticator() + + assert authenticator.auth_file == str(tmp_path / "env-dir" / "env.json") + + def test_bare_filename_auth_file_resolves_against_cwd( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.chdir(tmp_path) + + authenticator: Final = Authenticator(auth_file="auth.json") + + assert authenticator.auth_file == str(tmp_path / "auth.json") + assert authenticator.token_dir == str(tmp_path) + + def test_two_auth_files_are_isolated(self, tmp_path: Path) -> None: + file_a: Final = tmp_path / "account-a.json" + file_b: Final = tmp_path / "account-b.json" + _write_auth_record(file_a, "token-a", "acct-a") + _write_auth_record(file_b, "token-b", "acct-b") + + authenticator_a: Final = Authenticator(auth_file=str(file_a)) + authenticator_b: Final = Authenticator(auth_file=str(file_b)) + + assert authenticator_a.get_access_token() == "token-a" + assert authenticator_a.get_account_id() == "acct-a" + assert authenticator_b.get_access_token() == "token-b" + assert authenticator_b.get_account_id() == "acct-b" + + def test_get_cached_authenticator_reuses_instance_per_path(self, tmp_path: Path) -> None: + file_a: Final = tmp_path / "account-a.json" + file_b: Final = tmp_path / "account-b.json" + + authenticator_a: Final = get_cached_authenticator(str(file_a)) + + assert get_cached_authenticator(str(file_a)) is authenticator_a + assert get_cached_authenticator(str(file_b)) is not authenticator_a + + def test_get_chatgpt_auth_file(self) -> None: + assert get_chatgpt_auth_file(None) is None + assert get_chatgpt_auth_file({}) is None + assert get_chatgpt_auth_file({"chatgpt_auth_file": ""}) is None + assert get_chatgpt_auth_file({"chatgpt_auth_file": "/a/auth.json"}) == "/a/auth.json" + assert get_chatgpt_auth_file(GenericLiteLLMParams()) is None + assert get_chatgpt_auth_file(GenericLiteLLMParams(chatgpt_auth_file="/b/auth.json")) == "/b/auth.json" + + def test_relative_and_absolute_auth_file_share_cached_authenticator( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.chdir(tmp_path) + get_cached_authenticator.cache_clear() + auth_file: Final = tmp_path / "x" / "auth.json" + auth_file.parent.mkdir(parents=True) + auth_file.write_text("{}") + + relative_key: Final = get_chatgpt_auth_file({"chatgpt_auth_file": "x/auth.json"}) + absolute_key: Final = get_chatgpt_auth_file({"chatgpt_auth_file": str(auth_file)}) + assert relative_key is not None and absolute_key is not None + + assert get_cached_authenticator(relative_key) is get_cached_authenticator(absolute_key) + + +def test_refresh_keeps_saved_workspace(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("time.time", lambda: 1700000000) + auth_file: Final = tmp_path / "selected.json" + auth_file.write_text( + json.dumps( + { + "access_token": "expired-example-token", + "refresh_token": "example-refresh-token", + "expires_at": 1, + "account_id": "acct-selected-workspace", + } + ) + ) + token: Final = _make_jwt({"exp": 4102444800, "https://api.openai.com/auth": {"chatgpt_account_id": "acct-default"}}) + authenticator: Final = Authenticator(auth_file=str(auth_file)) + with respx.mock as transport: + transport.post("https://auth.openai.com/oauth/token").mock( + return_value=httpx.Response( + 200, + json={"access_token": token, "id_token": token, "refresh_token": "example-rotated-refresh"}, + ) + ) + assert authenticator.get_access_token() == token + assert authenticator.get_account_id(token) == "acct-selected-workspace" + assert json.loads(auth_file.read_text())["refresh_token"] == "example-rotated-refresh" diff --git a/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py new file mode 100644 index 00000000000..6b280d97089 --- /dev/null +++ b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py @@ -0,0 +1,90 @@ +import base64 +import json + +import pytest + +from litellm.llms.chatgpt.common_utils import ( + extract_chatgpt_account_id, + get_chatgpt_client_credential, + is_chatgpt_oauth_key, + select_chatgpt_client_credential_headers, +) + + +def _make_jwt(payload: dict) -> str: + def _b64(obj: dict) -> str: + raw = json.dumps(obj, separators=(",", ":")).encode("utf-8") + return base64.urlsafe_b64encode(raw).decode("utf-8").rstrip("=") + + return f"{_b64({'alg': 'none', 'typ': 'JWT'})}.{_b64(payload)}." + + +CHATGPT_JWT = _make_jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-alice"}}) +CHATGPT_JWT_WITHOUT_ACCOUNT = _make_jwt({"https://api.openai.com/auth": {"user_id": "user-1"}}) +NON_CHATGPT_JWT = _make_jwt({"sub": "someone", "https://api.openai.com/auth": "not-an-object"}) + + +class TestExtractChatGPTAccountId: + def test_reads_account_id_from_auth_claim(self): + assert extract_chatgpt_account_id(CHATGPT_JWT) == "acct-alice" + + @pytest.mark.parametrize( + "token", + [None, "", "sk-proj-not-a-jwt", "header.only", NON_CHATGPT_JWT, CHATGPT_JWT_WITHOUT_ACCOUNT], + ) + def test_returns_none_without_a_usable_account_claim(self, token): + assert extract_chatgpt_account_id(token) is None + + +class TestIsChatGPTOAuthKey: + def test_recognizes_bearer_and_raw_chatgpt_jwt(self): + assert is_chatgpt_oauth_key(CHATGPT_JWT) + assert is_chatgpt_oauth_key(f"Bearer {CHATGPT_JWT}") + assert is_chatgpt_oauth_key(f"Bearer {CHATGPT_JWT_WITHOUT_ACCOUNT}") + + @pytest.mark.parametrize( + "value", [None, "", "Bearer sk-proj-platform-key", "Bearer not.a.jwt", _make_jwt({"sub": "x"})] + ) + def test_rejects_non_chatgpt_credentials(self, value): + assert not is_chatgpt_oauth_key(value) + + +class TestSelectChatGPTClientCredentialHeaders: + def test_account_header_alone_is_not_forwarded(self): + selected = select_chatgpt_client_credential_headers( + {"chatgpt-account-id": "acct-alice", "authorization": "Bearer sk-proj-platform-key"} + ) + assert dict(selected) == {} + assert get_chatgpt_client_credential({"chatgpt-account-id": "acct-alice"}) is None + + def test_forwards_only_bearer_and_account_context(self): + selected = select_chatgpt_client_credential_headers( + { + "Authorization": f"Bearer {CHATGPT_JWT}", + "ChatGPT-Account-Id": "acct-explicit", + "x-litellm-api-key": "sk-gateway", + "user-agent": "codex", + } + ) + assert dict(selected) == {"Authorization": f"Bearer {CHATGPT_JWT}", "ChatGPT-Account-Id": "acct-explicit"} + credential = get_chatgpt_client_credential(selected) + assert credential is not None + assert credential.access_token == CHATGPT_JWT + assert credential.account_id == "acct-explicit" + + def test_account_id_falls_back_to_jwt_claim(self): + credential = get_chatgpt_client_credential({"authorization": f"Bearer {CHATGPT_JWT}"}) + assert credential is not None + assert credential.account_id == "acct-alice" + + +@pytest.mark.parametrize("scheme", ["Bearer", "bearer", "BEARER"]) +def test_bearer_scheme_is_case_insensitive(scheme): + credential = get_chatgpt_client_credential({"authorization": f"{scheme} {CHATGPT_JWT}"}) + assert credential is not None + assert credential.access_token == CHATGPT_JWT + assert credential.account_id == "acct-alice" + + +def test_other_authentication_schemes_are_not_chatgpt_logins(): + assert not is_chatgpt_oauth_key(f"Basic {CHATGPT_JWT}") diff --git a/tests/unit/llms/compactifai/test_compactifai.py b/tests/unit/llms/compactifai/test_compactifai.py index 1367c703fda..56d5374f383 100644 --- a/tests/unit/llms/compactifai/test_compactifai.py +++ b/tests/unit/llms/compactifai/test_compactifai.py @@ -2,7 +2,7 @@ import json import os import sys from unittest.mock import AsyncMock, patch -from typing import Optional +from typing import Final, Optional import httpx import pytest @@ -271,3 +271,38 @@ async def test_compactifai_async_completion(respx_mock): assert response.choices[0].message.content == "Async response from CompactifAI" assert response.usage.total_tokens == 23 + + +@pytest.mark.parametrize( + ("body", "json_mode", "error", "message"), + [ + ([], None, TypeError, r"argument after \*\* must be a mapping, not list"), + ({"choices": None}, True, TypeError, "'NoneType' object is not iterable"), + ({"choices": [{"message": "hi"}]}, True, AttributeError, "'str' object has no attribute 'get'"), + ({"choices": [{"message": {"tool_calls": [5]}}]}, True, TypeError, "'int' object is not subscriptable"), + ], +) +def test_compactifai_malformed_response_body_raises_the_error_of_the_failing_lookup(body, json_mode, error, message): + from unittest.mock import MagicMock + + from litellm.llms.compactifai.chat.transformation import CompactifAIChatConfig + + raw_response: Final = httpx.Response( + 200, json=body, request=httpx.Request("POST", "https://api.compactif.ai/v1/chat/completions") + ) + + with pytest.raises(error, match=message) as raised: + CompactifAIChatConfig().transform_response( + model="cai-llama-3-1-8b-slim", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + json_mode=json_mode, + ) + + assert raised.type is error diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index bd4ed4c99cc..1607447f7bf 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -2172,3 +2172,86 @@ async def test_async_ollama_ssl_verify_false_builds_a_client_that_skips_certific assert route.call_count == 1 assert text == "Sunny" assert _cached_clients_verify_modes(client_cache, AsyncHTTPHandler) == (ssl.CERT_NONE,) + + +_TIMEOUT_CLIENT_LIMITS: Final = httpx.Timeout(42.0, connect=7.0) +_TIMEOUT_PHASES: Final = ( + (httpx.ConnectTimeout, "Connection timed out", "Timeout passed=7.0"), + (httpx.ReadTimeout, "Read timed out", "Timeout passed=42.0"), + (httpx.WriteTimeout, "Write timed out", "Timeout passed=42.0"), + (httpx.PoolTimeout, "Connection pool timed out", "Timeout passed=42.0"), +) + + +def _timing_out(kind: type[httpx.TimeoutException]) -> httpx.MockTransport: + def _raise(request: httpx.Request) -> httpx.Response: + raise kind("simulated", request=request) + + return httpx.MockTransport(_raise) + + +@pytest.mark.parametrize("verb", ["post", "put", "patch", "delete"]) +@pytest.mark.parametrize(("kind", "label", "limit"), _TIMEOUT_PHASES) +def test_sync_timeout_without_a_per_call_timeout_names_the_phase_and_the_limit_that_fired( + verb: str, kind: type[httpx.TimeoutException], label: str, limit: str +) -> None: + """Issue #14635: the handler used to print the caller's absent argument as 'after None seconds'.""" + handler: Final = HTTPHandler() + handler.client.close() + handler.client = httpx.Client(transport=_timing_out(kind), timeout=_TIMEOUT_CLIENT_LIMITS) + try: + with pytest.raises(LitellmTimeout) as raised: + getattr(handler, verb)("https://example.com/v1/embeddings", json={"input": "hi"}) + finally: + handler.close() + + message: Final = str(raised.value) + assert label in message, message + assert limit in message, message + assert "None" not in message, message + + +@pytest.mark.parametrize("verb", ["post", "put", "patch"]) +@pytest.mark.parametrize(("kind", "label", "limit"), _TIMEOUT_PHASES) +async def test_async_timeout_without_a_per_call_timeout_names_the_phase_and_the_limit_that_fired( + verb: str, kind: type[httpx.TimeoutException], label: str, limit: str +) -> None: + handler: Final = AsyncHTTPHandler(timeout=_TIMEOUT_CLIENT_LIMITS) + await handler.client.aclose() + handler.client = httpx.AsyncClient(transport=_timing_out(kind), timeout=_TIMEOUT_CLIENT_LIMITS) + try: + with pytest.raises(LitellmTimeout) as raised: + await getattr(handler, verb)("https://example.com/v1/embeddings", json={"input": "hi"}) + finally: + await handler.close() + + message: Final = str(raised.value) + assert label in message, message + assert limit in message, message + assert "None" not in message, message + + +def test_sync_timeout_reports_the_per_call_timeout_as_a_float_when_one_is_given() -> None: + handler: Final = HTTPHandler() + handler.client.close() + handler.client = httpx.Client(transport=_timing_out(httpx.ReadTimeout), timeout=_TIMEOUT_CLIENT_LIMITS) + try: + with pytest.raises(LitellmTimeout) as raised: + handler.post("https://example.com/v1/embeddings", json={"input": "hi"}, timeout=30) + finally: + handler.close() + + assert "Read timed out. Timeout passed=30.0, time taken=" in str(raised.value), str(raised.value) + + +async def test_async_timeout_reports_the_per_call_timeout_as_a_float_when_one_is_given() -> None: + handler: Final = AsyncHTTPHandler(timeout=_TIMEOUT_CLIENT_LIMITS) + await handler.client.aclose() + handler.client = httpx.AsyncClient(transport=_timing_out(httpx.ReadTimeout), timeout=_TIMEOUT_CLIENT_LIMITS) + try: + with pytest.raises(LitellmTimeout) as raised: + await handler.post("https://example.com/v1/embeddings", json={"input": "hi"}, timeout=30) + finally: + await handler.close() + + assert "Read timed out. Timeout passed=30.0, time taken=" in str(raised.value), str(raised.value) 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 63cc45b29df..f7850aaf9a1 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -24,6 +24,7 @@ from litellm.integrations.code_interpreter_interception.handler import ( from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig +from litellm.llms.azure.vector_stores.transformation import AzureOpenAIVectorStoreConfig from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -5542,3 +5543,155 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_ 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) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("environment_base", [None, "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) +@pytest.mark.parametrize( + ("api_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/inference/go/anthropic", None, "conversation-123"), + ("https://opencode.ai/inference/go/anthropic", "caller-override", "caller-override"), + ("https://other.example", None, None), + ], +) +async def test_messages_opencode_session_header( + stream: bool, + environment_base: str | None, + api_base: str, + explicit_header: str | None, + expected_session: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + if environment_base: + monkeypatch.setenv(environment_base, "https://opencode.ai/inference/go/anthropic") + received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + payload: Final = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "OK"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + + def respond(request: httpx.Request) -> httpx.Response: + received.put_nowait(request) + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"event: message_start\ndata: {json.dumps({'type': 'message_start', 'message': payload})}\n\n" + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ) + if stream + else httpx.Response(200, json=payload) + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + result: Final = await litellm.anthropic.messages.acreate( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": "hello"}], + max_tokens=128000, + api_base=None if environment_base and api_base == "https://opencode.ai/inference/go/anthropic" else api_base, + api_key="test-key", + stream=stream, + litellm_session_id="conversation-123", + client=upstream, + extra_headers={"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + ) + if stream: + chunks: Final = b"".join([chunk async for chunk in result]) + assert b'"text": "OK"' in chunks + assert b"message_stop" in chunks + else: + assert result["content"] == [{"type": "text", "text": "OK"}] + finally: + await upstream.close() + sent: Final = received.get_nowait() + assert str(sent.url) == f"{api_base}/v1/messages" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["x-api-key"] == "test-key" + assert json.loads(sent.content) == { + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 128000, + "stream": stream, + } + assert received.empty() + + +def test_vector_store_retrieve_handler_sends_id_before_azure_api_version_query() -> None: + received: Final[list[httpx.Request]] = [] + + def answer(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response(200, json={"id": "vs_abc", "object": "vector_store"}) + + handler: Final = BaseLLMHTTPHandler() + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(answer))) + + handler.vector_store_retrieve_handler( + vector_store_id="vs_abc", + vector_store_provider_config=AzureOpenAIVectorStoreConfig(), + custom_llm_provider="azure", + litellm_params=GenericLiteLLMParams( + api_base="https://x.openai.azure.com", + api_key="azure-test-key", + api_version="2025-04-01-preview", + ), + logging_obj=Mock(), + client=client, + _is_async=False, + ) + + sent: Final = received[0] + assert str(sent.url) == "https://x.openai.azure.com/openai/vector_stores/vs_abc?api-version=2025-04-01-preview" + assert len(received) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.parametrize("provider", ("openai", "azure")) +async def test_non_chatgpt_responses_preserves_caller_headers_after_validation( + is_async: bool, provider: str, respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=r"https://responses\.example\.com/.*").respond( + 200, + json={ + "id": "resp_headers", "object": "response", "created_at": 1, + "model": "gpt-6-luna", "status": "completed", "output": [], + "parallel_tool_calls": False, "tool_choice": "auto", "tools": [], + }, + ) + client: Final = AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport()) if is_async else HTTPHandler(client=httpx.Client()) + kwargs: Final = { + "model": f"{provider}/gpt-6-luna", + "input": "hello", + "api_key": "fake-configured-key", + "api_base": "https://responses.example.com/v1", + "api_version": "2025-04-01-preview", + "extra_headers": {"Authorization": "Bearer fake-caller-key", "content-type": "application/vnd.example+json"}, + "client": client, + } + try: + response: Final = await litellm.aresponses(**kwargs) if is_async else litellm.responses(**kwargs) + assert response.status == "completed" + assert route.calls[0].request.headers["authorization"] == "Bearer fake-caller-key" + expected_content_types: Final = ( + ["application/json", "application/vnd.example+json"] + if provider == "openai" + else ["application/vnd.example+json"] + ) + assert route.calls[0].request.headers.get_list("content-type") == expected_content_types + finally: + if isinstance(client, AsyncHTTPHandler): + await client.close() + else: + client.client.close() diff --git a/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py b/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py index aff90549a7b..0951ba6ba92 100644 --- a/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py +++ b/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py @@ -1,6 +1,7 @@ import io import os import pathlib +from typing import Final from unittest.mock import MagicMock import httpx @@ -594,3 +595,24 @@ def test_transform_response_wraps_malformed_words_with_the_raw_body(words: objec DeepgramAudioTranscriptionConfig().transform_audio_transcription_response(raw_response) assert str(exc_info.value).endswith(f"\nResponse: {raw_response.text}") + + +@pytest.mark.parametrize( + ("payload", "error"), + [ + ({"results": None}, "'NoneType' object is not subscriptable"), + (_deepgram_payload({"transcript": "hello", "words": 5}, {}), "object of type 'int' has no len()"), + ( + {**_deepgram_payload({"transcript": "hello"}, {}), "metadata": None}, + "'NoneType' object is not subscriptable", + ), + ], +) +def test_transform_response_wraps_the_error_of_the_failing_lookup_in_a_malformed_body(payload: object, error: str): + raw_response: Final = httpx.Response(200, json=payload) + + with pytest.raises(ValueError, match="Error transforming Deepgram response: ") as exc_info: + DeepgramAudioTranscriptionConfig().transform_audio_transcription_response(raw_response) + + assert exc_info.type is ValueError + assert str(exc_info.value) == f"Error transforming Deepgram response: {error}\nResponse: {raw_response.text}" diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index dd99ad8bf91..3fbb12881da 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -511,6 +511,54 @@ class TestDeepSeekVisionMultimodalContent: assert any(block.get("type") == "image_url" for block in content) +@pytest.mark.usefixtures("_deepseek_httpx_transport") +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("model", ("deepseek-flash", "deepseek-v4-pro", "deepseek-chat", "deepseek-reasoner")) +@pytest.mark.parametrize("effort", ("minimal", "low", "medium", "high", "xhigh", "max", "ultra")) +@pytest.mark.parametrize("thinking", (None, "enabled")) +@pytest.mark.respx(assert_all_called=True) +async def test_deepseek_forwards_native_effort( + respx_mock: respx.MockRouter, model: str, effort: str, thinking: str | None, stream: bool +) -> None: + route: Final = respx_mock.post(DEEPSEEK_CHAT_COMPLETIONS_URL).mock(return_value=_deepseek_response(stream)) + response: Final = await litellm.acompletion( + model=f"deepseek/{model}", + messages=[{"role": "user", "content": "Hello"}], + api_key=DEEPSEEK_API_KEY, + reasoning_effort=effort, + stream=stream, + **({"thinking": {"type": thinking}} if thinking is not None else {}), + ) + if stream: + async for _ in response: + pass + body: Final = json.loads(route.calls[0].request.content) + assert body["reasoning_effort"] == effort + assert body["thinking"] == {"type": "enabled"} + + +@pytest.mark.parametrize( + "thinking,effort,expected_thinking", + ( + (None, None, None), + (None, "none", "disabled"), + ("enabled", "none", "enabled"), + ("disabled", "max", "disabled"), + ("enabled", None, "enabled"), + (None, "default", "enabled"), + (None, "bogus", "enabled"), + (None, {"effort": "high"}, "enabled"), + ), +) +def test_deepseek_preserves_toggle_without_forwarding_unsupported_effort( + thinking: str | None, effort: str | dict[str, str] | None, expected_thinking: str | None +) -> None: + params: Final = {"reasoning_effort": effort, **({"thinking": {"type": thinking}} if thinking is not None else {})} + result: Final = DeepSeekChatConfig().map_openai_params(params, {}, "deepseek-flash", False) + assert "reasoning_effort" not in result + assert result.get("thinking") == ({"type": expected_thinking} if expected_thinking else None) + + class TestDeepSeekThinkingParams: """Test thinking and reasoning_effort parameter handling for DeepSeek.""" @@ -735,9 +783,7 @@ async def test_deepseek_provider_async_completion(stream): messages = [{"role": "user", "content": "Hello, world!"}] # Mock AsyncHTTPHandler.post method for async test - with patch( - "litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post") as mock_post: mock_response_data = litellm.ModelResponse( choices=[ litellm.Choices( @@ -775,9 +821,7 @@ async def test_deepseek_provider_async_completion(stream): # Check request body request_body = json.loads(call_args.kwargs["data"]) assert call_args.kwargs["url"] == "https://api.deepseek.com/beta/chat/completions" - assert ( - request_body["model"] == "deepseek-reasoner" - ) # Model name should be stripped of provider prefix + assert request_body["model"] == "deepseek-reasoner" # Model name should be stripped of provider prefix assert request_body["messages"] == messages assert request_body["stream"] == stream diff --git a/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py b/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py index efdb428db32..ee6ee187d64 100644 --- a/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py +++ b/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py @@ -98,6 +98,12 @@ class TestEmbedding: class TestErrors: + def test_a_non_object_body_fails_as_an_invalid_response_object(self, eden_key, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(200, json=["unexpected"])) + + with pytest.raises(litellm.APIConnectionError, match="Invalid response object"): + litellm.embedding(model=MODEL, input="hello") + def test_middleware_401_maps_to_authentication_error(self, eden_key, respx_mock): respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token."})) diff --git a/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py b/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py index d7b32affc70..4e0673b8c27 100644 --- a/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py +++ b/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py @@ -117,6 +117,12 @@ class TestImageGeneration: class TestErrors: + def test_a_non_object_body_fails_as_an_invalid_response_object(self, eden_key, respx_mock): + respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(200, json=["unexpected"])) + + with pytest.raises(litellm.APIConnectionError, match="Invalid response object"): + litellm.image_generation(model=MODEL, prompt="a red square") + def test_middleware_401_maps_to_authentication_error(self, eden_key, respx_mock): respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token."})) diff --git a/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py b/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py index 86252a1a95d..f55644eccae 100644 --- a/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py +++ b/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py @@ -314,4 +314,4 @@ class TestElevenLabsTextToSpeechTransformation: ) assert voice_id in url - assert "output_format=pcm_44100" in url + assert "output_format=pcm_24000" in url diff --git a/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py index d05371d7df9..88f0ae1d524 100644 --- a/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py +++ b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py @@ -55,3 +55,31 @@ def test_speech_keeps_an_internal_prefixed_kwarg_out_of_the_elevenlabs_request(r sent: Final = json.loads(respx_mock.calls[0].request.content) assert "_litellm_undeclared_sentinel" not in sent, sent assert sent["text"] == "hi" + + +@pytest.mark.parametrize( + ("response_format", "expected_output_format"), + [("wav", "wav_24000"), ("pcm", "pcm_24000")], +) +def test_speech_maps_wav_and_pcm_to_elevenlabs_output_formats( + respx_mock: respx.MockRouter, + response_format: str, + expected_output_format: str, +) -> None: + api_base: Final = "http://localhost:12346" + mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/v1/text-to-speech/.*").mock( + return_value=httpx.Response(status_code=200, content=b"audio", headers={"content-type": "audio/mpeg"}) + ) + + litellm.speech( + model="elevenlabs/eleven_multilingual_v2", + input="hi", + voice="21m00Tcm4TlvDq8ikWAM", + response_format=response_format, + api_base=api_base, + api_key="fake_elevenlabs_api_key", + ) + + assert mock_route.called + request_url: Final = httpx.URL(str(respx_mock.calls[0].request.url)) + assert request_url.params["output_format"] == expected_output_format diff --git a/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py b/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py index da8ab138416..175ce040d9f 100644 --- a/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py +++ b/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py @@ -106,6 +106,26 @@ def test_generate_content_response_parses_function_call() -> None: } +@pytest.mark.parametrize( + ("body", "message"), + [ + ([], r"argument after \*\* must be a mapping, not list"), + (["candidates"], "list indices must be integers or slices, not str"), + ], +) +def test_generate_content_response_non_object_body_raises_the_error_of_the_failing_lookup( + body: JsonValue, message: str +) -> None: + with pytest.raises(TypeError, match=message) as raised: + GoogleGenAIConfig().transform_generate_content_response( + model="gemini-2.5-flash", + raw_response=httpx.Response(200, json=body), + logging_obj=SimpleNamespace(model_call_details={}), + ) + + assert raised.type is TypeError + + def test_generate_content_request_preserves_tool_declarations() -> None: tools: Final = [ { diff --git a/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index bd9b7006e58..c21dec9b6bb 100644 --- a/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -241,6 +241,26 @@ class TestGeminiImageEditTransformation: ) assert logging_usage["completion_tokens_details"]["image_tokens"] == 1716 + @pytest.mark.parametrize( + ("body", "error", "message"), + [ + (["not", "an", "object"], AttributeError, "'list' object has no attribute 'get'"), + ({"candidates": None}, TypeError, "'NoneType' object is not iterable"), + ({"candidates": [{"content": {"parts": None}}]}, TypeError, "'NoneType' object is not iterable"), + ], + ) + def test_transform_image_edit_response_malformed_body_raises_the_error_of_the_failing_lookup( + self, body: object, error: type[Exception], message: str + ) -> None: + with pytest.raises(error, match=message) as raised: + self.config.transform_image_edit_response( + model=self.model, + raw_response=httpx.Response(200, json=body), + logging_obj=self.logging_obj, + ) + + assert raised.type is error + def test_transform_image_edit_request_without_image_raises(self) -> None: optional_params = {} diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 196b912752b..20c36878fef 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -355,6 +355,25 @@ class TestGeminiModelInfo: assert result == expected +@pytest.mark.parametrize( + ("body", "error", "message"), + [ + ([], TypeError, "list indices must be integers or slices, not str"), + ({"models": None}, TypeError, "'NoneType' object is not iterable"), + ({"models": [{"name": 1}]}, AttributeError, "'int' object has no attribute 'replace'"), + ], +) +def test_get_models_malformed_listing_raises_the_error_of_the_failing_lookup( + body: JsonValue, error: type[Exception], message: str, respx_mock: respx.MockRouter +) -> None: + respx_mock.get("https://models.test/v1beta/models").respond(200, json=body) + + with pytest.raises(error, match=message) as raised: + GeminiModelInfo().get_models(api_key="test-key", api_base="https://models.test") + + assert raised.type is error + + class TestGoogleAIStudioTokenCounter: """Test suite for GoogleAIStudioTokenCounter class""" diff --git a/tests/unit/llms/gemini/test_gemini_tts.py b/tests/unit/llms/gemini/test_gemini_tts.py index a7522f6b315..d0ef04e6372 100644 --- a/tests/unit/llms/gemini/test_gemini_tts.py +++ b/tests/unit/llms/gemini/test_gemini_tts.py @@ -2,14 +2,70 @@ Test Gemini TTS (Text-to-Speech) functionality """ -import pytest +from typing import Final from unittest.mock import patch, MagicMock +import pytest import litellm +from litellm.endpoints.speech.speech_to_completion_bridge.transformation import ( + SpeechToCompletionBridgeTransformationHandler, +) from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.utils import get_supported_openai_params +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + +GEMINI_3_1_FLASH_TTS_MODEL = "gemini-3.1-flash-tts-preview" + +MULTI_SPEAKER_SPEECH_CONFIG = { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + { + "speaker": "Ryan", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Umbriel", + }, + }, + }, + { + "speaker": "Katie", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Leda", + }, + }, + }, + ], + }, + "language_code": "en-US", +} + +NORMALIZED_MULTI_SPEAKER_SPEECH_CONFIG = { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speaker": "Ryan", + "voiceConfig": { + "prebuiltVoiceConfig": { + "voiceName": "Umbriel", + }, + }, + }, + { + "speaker": "Katie", + "voiceConfig": { + "prebuiltVoiceConfig": { + "voiceName": "Leda", + }, + }, + }, + ], + }, + "languageCode": "en-US", +} + class TestGeminiTTSTransformation: """Test Gemini TTS transformation functionality""" @@ -19,17 +75,16 @@ class TestGeminiTTSTransformation: config = GoogleAIStudioGeminiConfig() # Test TTS models (both preview and non-preview versions) - assert ( - config.is_model_gemini_audio_model("gemini-2.5-flash-preview-tts") == True - ) - assert config.is_model_gemini_audio_model("gemini-2.5-pro-preview-tts") == True - assert config.is_model_gemini_audio_model("gemini-2.5-flash-tts") == True - assert config.is_model_gemini_audio_model("gemini-2.5-pro-tts") == True + assert config.is_model_gemini_audio_model("gemini-2.5-flash-preview-tts") + assert config.is_model_gemini_audio_model("gemini-2.5-pro-preview-tts") + assert config.is_model_gemini_audio_model("gemini-2.5-flash-tts") + assert config.is_model_gemini_audio_model("gemini-2.5-pro-tts") + assert config.is_model_gemini_audio_model(GEMINI_3_1_FLASH_TTS_MODEL) # Test non-TTS models - assert config.is_model_gemini_audio_model("gemini-2.5-flash") == False - assert config.is_model_gemini_audio_model("gemini-2.5-pro") == False - assert config.is_model_gemini_audio_model("gpt-4o-audio-preview") == False + assert not config.is_model_gemini_audio_model("gemini-2.5-flash") + assert not config.is_model_gemini_audio_model("gemini-2.5-pro") + assert not config.is_model_gemini_audio_model("gpt-4o-audio-preview") def test_gemini_tts_supported_params(self): """Test that audio parameter is included for TTS models""" @@ -66,10 +121,7 @@ class TestGeminiTTSTransformation: assert "speechConfig" in result assert "voiceConfig" in result["speechConfig"] assert "prebuiltVoiceConfig" in result["speechConfig"]["voiceConfig"] - assert ( - result["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] - == "Kore" - ) + assert result["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" # Check response modalities assert "responseModalities" in result @@ -78,9 +130,7 @@ class TestGeminiTTSTransformation: def test_gemini_tts_audio_parameter_mapping_with_language_code(self): config = GoogleAIStudioGeminiConfig() - non_default_params = { - "audio": {"voice": "Kore", "format": "pcm16", "language_code": "en-US"} - } + non_default_params = {"audio": {"voice": "Kore", "format": "pcm16", "language_code": "en-US"}} optional_params = {} result = config.map_openai_params( @@ -102,6 +152,16 @@ class TestGeminiTTSTransformation: assert result["languageCode"] == "de-DE" assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + def test_map_audio_params_keeps_voice_beside_nested_speech_config(self): + config = GoogleAIStudioGeminiConfig() + + result = config._map_audio_params( + {"speech_config": {"language_code": "fr-FR"}, "voice": "Kore", "format": "pcm16"} + ) + + assert result["languageCode"] == "fr-FR" + assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + def test_map_audio_params_no_language_code(self): config = GoogleAIStudioGeminiConfig() @@ -110,6 +170,55 @@ class TestGeminiTTSTransformation: assert "languageCode" not in result assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + def test_map_audio_params_keeps_nested_speaker_list_and_drops_unknown_fields(self): + config: Final = GoogleAIStudioGeminiConfig() + + result: Final = config._map_audio_params( + { + "speech_config": { + "language_code": "fr-FR", + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Ryan", "voice_config": {"prebuilt_voice_config": {"voice_name": "Umbriel"}}} + ] + }, + "metadata": ["studio", 2], + }, + "format": "pcm16", + } + ) + + assert result == { + "languageCode": "fr-FR", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Umbriel"}}} + ] + }, + } + + def test_gemini_tts_multi_speaker_audio_parameter_mapping(self): + """Test multi-speaker audio parameter mapping for Gemini 3.1 TTS models""" + config = GoogleAIStudioGeminiConfig() + + non_default_params = { + "audio": { + "speech_config": MULTI_SPEAKER_SPEECH_CONFIG, + "format": "pcm16", + } + } + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=GEMINI_3_1_FLASH_TTS_MODEL, + drop_params=False, + ) + + assert result["speechConfig"] == NORMALIZED_MULTI_SPEAKER_SPEECH_CONFIG + assert result["responseModalities"] == ["AUDIO"] + def test_gemini_tts_audio_parameter_with_existing_modalities(self): """Test audio parameter mapping when modalities already exist""" config = GoogleAIStudioGeminiConfig() @@ -193,9 +302,7 @@ class TestGeminiTTSTransformation: } optional_params = {} - with pytest.raises( - ValueError, match="Unsupported audio format for Gemini TTS models" - ): + with pytest.raises(ValueError, match="Unsupported audio format for Gemini TTS models"): config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -234,6 +341,28 @@ def test_gemini_tts_completion_mock(): assert response.choices[0].message.content is not None +def test_gemini_tts_speech_bridge_accepts_multi_speaker_voice_dict(): + handler = SpeechToCompletionBridgeTransformationHandler() + + result = handler.transform_request( + model=f"vertex_ai/{GEMINI_3_1_FLASH_TTS_MODEL}", + input="Ryan: How are you doing today Katie?\nKatie: Not too bad.", + voice=MULTI_SPEAKER_SPEECH_CONFIG, + optional_params={"response_format": "pcm"}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + ) + + assert result["modalities"] == ["audio"] + assert result["audio"] == { + "speech_config": NORMALIZED_MULTI_SPEAKER_SPEECH_CONFIG, + "format": "pcm16", + } + assert "response_format" not in result + + class TestGeminiTTSSpeechConfigInRequestBody: """Test that speechConfig is properly included in the final request body. @@ -250,6 +379,8 @@ class TestGeminiTTSSpeechConfigInRequestBody: ("gemini-2.5-flash-preview-tts", "vertex_ai"), ("gemini-2.5-flash-preview-tts", "gemini"), ("gemini-2.5-pro-tts", "vertex_ai"), + (GEMINI_3_1_FLASH_TTS_MODEL, "vertex_ai"), + (GEMINI_3_1_FLASH_TTS_MODEL, "gemini"), ], ) def test_speechconfig_in_generation_config_transform_request_body(self, model, custom_llm_provider): @@ -283,12 +414,7 @@ class TestGeminiTTSSpeechConfigInRequestBody: f"speechConfig was filtered out of generationConfig for model={model}, provider={custom_llm_provider}. " "Ensure speechConfig is in the GenerationConfig TypedDict." ) - assert ( - generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"][ - "voiceName" - ] - == "Kore" - ) + assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" @pytest.mark.parametrize( "model,custom_llm_provider", @@ -296,6 +422,8 @@ class TestGeminiTTSSpeechConfigInRequestBody: ("gemini-2.5-flash-tts", "vertex_ai"), ("gemini-2.5-flash-tts", "gemini"), ("gemini-2.5-flash-preview-tts", "vertex_ai"), + (GEMINI_3_1_FLASH_TTS_MODEL, "vertex_ai"), + (GEMINI_3_1_FLASH_TTS_MODEL, "gemini"), ], ) def test_speechconfig_end_to_end_mapping(self, model, custom_llm_provider): @@ -342,17 +470,54 @@ class TestGeminiTTSSpeechConfigInRequestBody: f"speechConfig was filtered out during transform_request_body() for model={model}, provider={custom_llm_provider}. " "This breaks Gemini TTS - speechConfig must be in GenerationConfig TypedDict." ) - assert ( - generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"][ - "voiceName" - ] - == "Puck" - ) + assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Puck" # Also verify responseModalities is present assert "responseModalities" in generation_config assert "AUDIO" in generation_config["responseModalities"] + @pytest.mark.parametrize( + "custom_llm_provider", + [ + "vertex_ai", + "gemini", + ], + ) + def test_multi_speaker_speechconfig_end_to_end_mapping(self, custom_llm_provider): + """Test full pipeline for Gemini 3.1 multi-speaker TTS speechConfig""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _transform_request_body, + ) + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + config = VertexGeminiConfig() + + mapped_params = config.map_openai_params( + non_default_params={ + "audio": { + "speech_config": MULTI_SPEAKER_SPEECH_CONFIG, + "format": "pcm16", + } + }, + optional_params={}, + model=GEMINI_3_1_FLASH_TTS_MODEL, + drop_params=False, + ) + + request_body = _transform_request_body( + messages=[{"role": "user", "content": "Ryan: Hi.\nKatie: Hello."}], + model=GEMINI_3_1_FLASH_TTS_MODEL, + optional_params=mapped_params, + custom_llm_provider=custom_llm_provider, + litellm_params={}, + cached_content=None, + ) + + generation_config = request_body["generationConfig"] + assert generation_config["speechConfig"] == NORMALIZED_MULTI_SPEAKER_SPEECH_CONFIG + assert generation_config["responseModalities"] == ["AUDIO"] @pytest.mark.parametrize( "model,custom_llm_provider", @@ -395,14 +560,35 @@ class TestGeminiTTSSpeechConfigInRequestBody: generation_config = request_body["generationConfig"] assert generation_config["speechConfig"]["languageCode"] == "pt-BR" - assert ( - generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"][ - "voiceName" - ] - == "Puck" - ) + assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Puck" assert "AUDIO" in generation_config["responseModalities"] if __name__ == "__main__": pytest.main([__file__]) + + +@pytest.mark.parametrize("voice_key", ["name", "voiceName", "voice_name", "voice"]) +@pytest.mark.parametrize("wrapper", [None, "speechConfig", "speech_config"]) +@pytest.mark.parametrize("model", ["gemini-2.5-flash-tts", "gemini-2.5-pro-tts", GEMINI_3_1_FLASH_TTS_MODEL]) +def test_speech_bridge_preserves_single_speaker_voice_mapping(voice_key: str, wrapper: str | None, model: str): + handler = SpeechToCompletionBridgeTransformationHandler() + config = {voice_key: "Umbriel", "language_code": "en-US"} + voice = {wrapper: config} if wrapper else config + result = handler.transform_request( + model=model, + input="Hello", + voice=voice, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + ) + assert result["audio"] == { + "speech_config": { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Umbriel"}}, + "languageCode": "en-US", + }, + "format": "pcm16", + } diff --git a/tests/unit/llms/groq/chat/test_groq_chat_transformation.py b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py index 95611e344dc..7aa8c65215e 100644 --- a/tests/unit/llms/groq/chat/test_groq_chat_transformation.py +++ b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py @@ -1,4 +1,5 @@ import logging +from typing import Final from unittest.mock import patch import httpx @@ -523,3 +524,68 @@ def test_tool_search_reply_sent_to_groq_drops_litellm_response_fields(): } ], } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("with_tool_call", [False, True]) +async def test_streamed_reasoning_reply_can_be_replayed(is_async: bool, with_tool_call: bool) -> None: + tool_calls: Final = [ + { + "index": 0, + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + } + ] + handler: Final = GroqChatCompletionStreamingHandler(streaming_response=iter(()), sync_stream=True) + chunk: Final = handler.chunk_parser( + { + "id": "chatcmpl-reasoning", + "object": "chat.completion.chunk", + "created": 1769511767, + "model": "openai/gpt-oss-20b", + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "reasoning": "I should check the weather.", + "content": "" if with_tool_call else "Hello there!", + **({"tool_calls": tool_calls} if with_tool_call else {}), + }, + "finish_reason": "tool_calls" if with_tool_call else "stop", + } + ], + } + ) + reply: Final = litellm.stream_chunk_builder([chunk]).choices[0].message + assert reply.reasoning_content == "I should check the weather." + history: Final = [reply.model_dump(exclude_none=True)] + config: Final = GroqChatConfig() + + if is_async: + replayed = await config._transform_messages(messages=history, model="openai/gpt-oss-20b", is_async=True) + else: + replayed = config._transform_messages(messages=history, model="openai/gpt-oss-20b", is_async=False) + + assert replayed == [ + { + "role": "assistant", + "content": "" if with_tool_call else "Hello there!", + **( + { + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + } + ] + } + if with_tool_call + else {} + ), + } + ] + assert reply.reasoning_content == "I should check the weather." diff --git a/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py b/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py index 2f1b3764610..04c2ce153c6 100644 --- a/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py +++ b/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py @@ -10,6 +10,7 @@ import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.huggingface.common_utils import _fetch_inference_provider_mapping +from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, _load_huggingface_tokenizer HUGGINGFACE_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct" HUGGINGFACE_ROUTER_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct-Turbo" @@ -83,9 +84,13 @@ def isolate_huggingface_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator monkeypatch.delenv("HF_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_KEY", raising=False) + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() _fetch_inference_provider_mapping.cache_clear() yield _fetch_inference_provider_mapping.cache_clear() + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() PROVIDER_MAPPING_RESPONSE = { diff --git a/tests/unit/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py index f465bbea663..9e2754eef68 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -364,7 +364,8 @@ def test_text_completion_stream_forwards_stream_options_and_usage(monkeypatch): ) ) - assert tuple(chunk.choices[0].text for chunk in chunks) == ("hello", None, None) + assert tuple(chunk.choices[0].text for chunk in chunks[:-1]) == ("hello", None) + assert chunks[-1].choices == [] assert tuple(chunk.usage for chunk in chunks[:-1]) == (None, None) assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == ( 2, diff --git a/tests/unit/llms/openai/evals/test_openai_evals_transformation.py b/tests/unit/llms/openai/evals/test_openai_evals_transformation.py index f359f214f97..ed8f86d4727 100644 --- a/tests/unit/llms/openai/evals/test_openai_evals_transformation.py +++ b/tests/unit/llms/openai/evals/test_openai_evals_transformation.py @@ -264,6 +264,17 @@ def test_transform_cancel_eval_request(config: OpenAIEvalsConfig): assert request_body == {} +def test_transform_cancel_eval_request_appends_cancel_to_an_eval_id_named_cancel(config: OpenAIEvalsConfig) -> None: + eval_url, _ = config.transform_get_eval_request( + eval_id="cancel", api_base="https://api.openai.com", litellm_params=GenericLiteLLMParams(), headers={} + ) + cancel_url, _, _ = config.transform_cancel_eval_request( + eval_id="cancel", api_base="https://api.openai.com", litellm_params=GenericLiteLLMParams(), headers={} + ) + + assert cancel_url == f"{eval_url}/cancel" + + def test_transform_cancel_eval_response(config: OpenAIEvalsConfig): """Test transformation of cancel eval response""" response = httpx.Response( diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index 0a5039791ee..8995b3cb068 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -8,6 +8,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.transformation impor from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder def test_transform_basic_request(): @@ -207,6 +208,23 @@ def test_messages_to_responses_input_with_tool(): } +def test_messages_to_responses_input_reads_a_tool_call_that_forwards_its_fields(): + tool_call: Final = {"id": "call_123", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + messages: Final = [ + {"role": "user", "content": "What's the weather?"}, + {"role": "assistant", "content": None, "tool_calls": [AttributeForwarder(tool_call)]}, + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items[1] == { + "type": "function_call", + "call_id": "call_123", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + } + + def test_messages_to_responses_input_preserves_images(): """An image block must survive the round trip, or OpenAI counts only the text. diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index bf871210074..0d1b1aca9bc 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -34,6 +34,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI from litellm.types.llms.openai import ChatCompletionToolCallChunk +from litellm.responses.litellm_completion_transformation.compaction_items import compaction_summary from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -46,6 +47,7 @@ from litellm.types.responses.main import ( OutputToolSearchCall, ) from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder class MockGuardrail(CustomGuardrail): @@ -298,6 +300,32 @@ class TestOpenAIResponsesHandlerInputProcessing: assert result["input"][1]["content"] == "World [GUARDRAILED]" assert result["model"] == "gpt-4" + @pytest.mark.asyncio + async def test_compaction_summary_is_scanned_and_rewritten_inside_its_item(self) -> None: + affinity_prefix: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwLWE=;" + data: Final = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": affinity_prefix + + json.dumps({"type": "compaction", "content": "Cat is Biscuit."}), + }, + {"role": "user", "content": "Cat name?", "type": "message"}, + ], + "model": "claude-sonnet-5-5", + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, MockGuardrail(guardrail_name="t")) + + rewritten: Final = result["input"][0]["encrypted_content"] + assert rewritten.startswith(affinity_prefix) + assert json.loads(rewritten.removeprefix(affinity_prefix)) == { + "type": "compaction", + "content": "Cat is Biscuit. [GUARDRAILED]", + } + assert result["input"][1]["content"] == "Cat name? [GUARDRAILED]" + @pytest.mark.asyncio async def test_process_input_list_with_multimodal_content(self): """Test processing list input with multimodal content""" @@ -1756,6 +1784,22 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_events_that_forward_their_fields(self): + handler: Final = OpenAIResponsesHandler() + plain_events: Final = self._ended_stream_events() + forwarded_delta: Final = AttributeForwarder(plain_events[0]) + events: Final = [forwarded_delta, *plain_events[1:]] + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert forwarded_delta.delta == "hello [MASKED]" + @staticmethod def _ended_function_call_stream_events() -> List[dict]: def item(arguments: str, status: str) -> dict: @@ -3072,6 +3116,21 @@ class DictStructuredMessagesGuardrail(CustomGuardrail): return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}} +class NameMaskingStructuredGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + rewritten: Final = [ + {**m, "content": m["content"].replace("Biscuit", "[NAME]")} if isinstance(m.get("content"), str) else m + for m in inputs.get("structured_messages") or [] + ] + return {**inputs, "structured_messages": rewritten} + + def _parallel_tool_call_input() -> list: return [ {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, @@ -3298,6 +3357,30 @@ class TestProvenancePatching: assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + @pytest.mark.asyncio + async def test_masked_compaction_summary_stays_a_compaction_item_that_drops_older_history(self) -> None: + data: Final = { + "model": "claude-sonnet-5-5", + "input": [ + {"role": "user", "content": "My cat is Biscuit. " * 200}, + {"role": "assistant", "content": "Noted."}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Cat is Biscuit."}), + }, + {"role": "user", "content": "Cat name?"}, + ], + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, NameMaskingStructuredGuardrail()) + + assert compaction_summary(result["input"][2]) == "Cat is [NAME]." + provider_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=result["input"], responses_api_request={}, only_since_compaction=True + ) + assert [message["content"] for message in provider_messages] == ["Cat is [NAME].", "Cat name?"] + def test_provenance_bails_on_non_mapping_item(self): from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance @@ -4174,3 +4257,117 @@ class TestOpenAIResponsesHandlerStreamingScanKey: assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hi"]], guardrail.seen_inputs tool_calls = guardrail.seen_inputs[0].get("tool_calls") or [] assert [call["function"]["arguments"] for call in tool_calls] == ['{"city": "Paris"}'], tool_calls + + @pytest.mark.asyncio + async def test_a_tool_call_finished_mid_stream_reaches_the_guardrail_without_texts(self): + guardrail: Final = MockRecordingGuardrail(guardrail_name="test") + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[self._finished_function_call(1, "fc_1", "Paris")[-1]], + guardrail_to_apply=guardrail, + request_data={}, + ) + assert [ + (sorted(inputs), [call["function"]["arguments"] for call in inputs.get("tool_calls") or []]) + for inputs in guardrail.seen_inputs + ] == [(["tool_calls"], ['{"city": "Paris"}'])] + + +class TestCompactionOutputSummaryReachesOutputGuardrails: + SUMMARY: Final = "The user lives at 1 Secret Lane." + + @classmethod + def _compaction_item(cls, summary: str | None = None) -> dict[str, str]: + return { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary or cls.SUMMARY}), + } + + @staticmethod + def _summary(item: object) -> str: + encrypted_content: Final = item["encrypted_content"] if isinstance(item, dict) else item.encrypted_content + return json.loads(encrypted_content)["content"] + + @staticmethod + def _lane_masking_guardrail() -> CustomGuardrail: + class MaskLane(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return { + **inputs, + "texts": [text.replace("Secret Lane", "[MASKED]") for text in inputs.get("texts", [])], + } + + return MaskLane(guardrail_name="mask-lane") + + @pytest.mark.asyncio + async def test_non_streaming_summary_is_masked_inside_its_item(self) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[self._compaction_item()], + ) + + result: Final = await OpenAIResponsesHandler().process_output_response(response, self._lane_masking_guardrail()) + + assert self._summary(result.output[0]) == "The user lives at 1 [MASKED]." + + @pytest.mark.asyncio + async def test_ended_stream_summary_is_masked_in_the_done_item_and_the_envelope(self) -> None: + events: Final = [ + { + "type": "response.output_item.added", + "output_index": 0, + "item": {**self._compaction_item(), "encrypted_content": ""}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": self._compaction_item()}, + { + "type": "response.completed", + "response": { + "id": "resp_1", + "created_at": 1, + "model": "claude-sonnet-5-5", + "output": [self._compaction_item()], + "status": "completed", + }, + }, + ] + + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._lane_masking_guardrail(), + deliver_ended_stream_rewrites=True, + ) + + assert self._summary(events[1]["item"]) == "The user lives at 1 [MASKED]." + assert self._summary(events[2]["response"]["output"][0]) == "The user lives at 1 [MASKED]." + + @pytest.mark.asyncio + async def test_a_blocking_guardrail_rejects_a_finished_compaction_item_mid_stream(self) -> None: + with pytest.raises(HTTPException) as exc_info: + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[ + {"type": "response.output_item.done", "output_index": 0, "item": self._compaction_item()} + ], + guardrail_to_apply=MockGuardrail(guardrail_name="block"), + ) + + assert exc_info.value.detail["texts"] == [self.SUMMARY] + + def test_a_changed_summary_changes_the_ended_stream_scan_key(self) -> None: + handler: Final = OpenAIResponsesHandler() + + def ended_key(summary: str) -> object: + return handler.get_streaming_scan_key( + [{"type": "response.completed", "response": {"output": [self._compaction_item(summary)]}}] + ) + + assert ended_key(self.SUMMARY) != ended_key("The user lives elsewhere.") diff --git a/tests/unit/llms/opencode/chat/__init__.py b/tests/unit/llms/opencode/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/opencode/chat/test_transformation.py b/tests/unit/llms/opencode/chat/test_transformation.py new file mode 100644 index 00000000000..0a5e49ca1b2 --- /dev/null +++ b/tests/unit/llms/opencode/chat/test_transformation.py @@ -0,0 +1,582 @@ +import asyncio +import json +from collections.abc import Mapping +from datetime import datetime +from typing import Final + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.opencode.common_utils import go_headers + + +@pytest.fixture +def go_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + for route in ("chat/completions", "messages", "responses", "gemini"): + monkeypatch.setitem( + litellm.model_cost, + f"opencode_go/fixture-{route.split('/')[0]}", + { + "litellm_provider": "opencode_go", + "mode": "responses" if route == "responses" else "chat", + "supported_endpoints": ["/v1beta/models" if route == "gemini" else f"/v1/{route}"], + "max_output_tokens": 128000, + "max_input_tokens": 1000000, + "input_cost_per_token": 0.000002, + "output_cost_per_token": 0.000005, + "supports_function_calling": True, + }, + ) + monkeypatch.delenv("OPENCODE_GO_API_BASE", raising=False) + + +def upstream_reply(route: str, model: str) -> Mapping[str, object]: + if route == "gemini": + return { + "candidates": [{"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, + } + if route == "messages": + return { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + if route == "responses": + return { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "model": model, + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_test", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + return { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1700000000, + "model": model, + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hello"}}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("upstream", ["chat/completions", "messages", "responses", "gemini"]) +@pytest.mark.parametrize("ingress", ["chat", "responses", "messages"]) +async def test_all_routes_use_model_protocol_and_preserve_session( + go_catalog: None, upstream: str, ingress: str +) -> None: + model: Final = f"fixture-{upstream.split('/')[0]}" + + def respond(request: httpx.Request) -> httpx.Response: + family: Final = "anthropic" if upstream == "messages" else "openai" + expected_path: Final = ( + f"/inference/go/google/v1beta/models/{model}:generateContent" + if upstream == "gemini" + else f"/inference/go/{family}/v1/{upstream}" + ) + assert request.url.path == expected_path + assert request.headers["authorization"] == "Bearer fake-go-key" + assert request.headers["x-opencode-session"] == "conversation-1" + body: Final = json.loads(request.content) + if upstream == "gemini": + assert body["contents"][0]["parts"][0]["text"] == "hello" + else: + assert body["model"] == model + assert "input" in body if upstream == "responses" else "messages" in body + if upstream == "messages": + assert body["max_tokens"] == 128000 + return httpx.Response(200, json=upstream_reply(upstream, model)) + + with respx.mock(using="httpx") as transport: + transport.post().mock(side_effect=respond) + kwargs: Final = { + "model": f"opencode_go/{model}", + "api_key": "fake-go-key", + "litellm_session_id": "conversation-1", + } + if ingress == "chat": + response: Final = await litellm.acompletion(messages=[{"role": "user", "content": "hello"}], **kwargs) + assert response.choices[0].message.content == "hello" + assert response.usage.total_tokens == 15 + elif ingress == "responses": + response_api: Final = await litellm.aresponses(input="hello", **kwargs) + assert response_api.output[0].content[0].text == "hello" + assert response_api.usage.total_tokens == 15 + else: + message: Final = await litellm.anthropic.messages.acreate( + messages=[{"role": "user", "content": "hello"}], + max_tokens=128000, + **kwargs, + ) + assert message["content"][0]["text"] == "hello" + assert message["usage"]["output_tokens"] == 5 + + +def test_go_never_falls_back_to_other_provider_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("OPENCODE_API_KEY", raising=False) + monkeypatch.delenv("OPENCODE_GO_API_KEY", raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "fake-openai-key") + monkeypatch.setenv("ANTHROPIC_API_KEY", "fake-anthropic-key") + with pytest.raises(litellm.AuthenticationError, match="OPENCODE_API_KEY"): + go_headers({}, {}, None, "fixture-model") + + +@pytest.mark.parametrize("input_tokens,expected", [(100, 0.00025), (101, 0.000504)]) +def test_go_uses_registered_context_tiers_not_underlying_provider_prices( + monkeypatch: pytest.MonkeyPatch, + input_tokens: int, + expected: float, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "opencode_go/fixture-priced", + { + "litellm_provider": "opencode_go", + "mode": "chat", + "input_cost_per_token": 0.000002, + "output_cost_per_token": 0.000005, + "tiered_pricing": [ + {"range": [0, 100], "input_cost_per_token": 0.000002, "output_cost_per_token": 0.000005}, + {"range": [100, 1000], "input_cost_per_token": 0.000004, "output_cost_per_token": 0.000010}, + ], + }, + ) + cost: Final = litellm.cost_per_token( + model="opencode_go/fixture-priced", + prompt_tokens=input_tokens, + completion_tokens=10, + ) + assert sum(cost) == pytest.approx(expected) + + +@pytest.mark.parametrize("suffix", ["", "/v1", "/v1/messages", "/v1/messages/"]) +def test_messages_custom_base_does_not_duplicate_version_or_route(go_catalog: None, suffix: str) -> None: + from litellm.llms.opencode.messages.transformation import OpenCodeGoMessagesConfig + + assert ( + OpenCodeGoMessagesConfig().get_complete_url( + f"https://gateway.example/go{suffix}", + None, + "fixture-messages", + {}, + {}, + ) + == "https://gateway.example/go/v1/messages" + ) + + +@pytest.mark.parametrize("base", [None, "https://gateway.example/go/v1"]) +def test_discovery_returns_qualified_ids_and_accepts_response_metadata(base: str | None) -> None: + url: Final = "https://gateway.example/go/v1/models" if base else "https://opencode.ai/inference/go/v1/models" + with respx.mock(using="httpx") as transport: + transport.get(url).respond( + json={ + "object": "list", + "data": [ + {"id": "fixture-chat", "object": "model"}, + {"id": "fixture-messages", "object": "model"}, + ], + } + ) + assert litellm.get_valid_models( + check_provider_endpoint=True, + custom_llm_provider="opencode_go", + api_key="fake-go-key", + api_base=base, + ) == [ + "opencode_go/fixture-chat", + "opencode_go/fixture-messages", + ] + + +@pytest.mark.parametrize("explicit_base", [None, "https://explicit-gateway.example/go/v1"]) +def test_discovery_honors_environment_base_and_explicit_override( + monkeypatch: pytest.MonkeyPatch, explicit_base: str | None +) -> None: + from litellm.llms.opencode.chat.transformation import OpenCodeGoChatConfig + + monkeypatch.setenv("OPENCODE_GO_API_BASE", "https://env-gateway.example/go/v1/") + expected_base: Final = explicit_base or "https://env-gateway.example/go/v1" + with respx.mock(using="httpx") as transport: + transport.get(f"{expected_base}/models").respond(json={"data": [{"id": "gateway-only-model"}]}) + assert OpenCodeGoChatConfig().get_models(api_key="fake-go-key", api_base=explicit_base) == [ + "opencode_go/gateway-only-model" + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ingress", ["chat", "responses", "messages"]) +async def test_gemini_streams_text_and_usage_through_all_routes(go_catalog: None, ingress: str) -> None: + def respond(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/inference/go/google/v1beta/models/fixture-gemini:streamGenerateContent" + assert request.url.query == b"alt=sse" + assert "stream" not in json.loads(request.content) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="data: " + json.dumps(upstream_reply("gemini", "fixture-gemini")) + "\n\n", + ) + + with respx.mock(using="httpx") as transport: + transport.post().mock(side_effect=respond) + kwargs: Final = {"model": "opencode_go/fixture-gemini", "api_key": "fake-go-key", "stream": True} + if ingress == "chat": + stream: Final = await litellm.acompletion( + messages=[{"role": "user", "content": "hello"}], stream_options={"include_usage": True}, **kwargs + ) + chunks: Final = [chunk async for chunk in stream] + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "hello" + assert chunks[-1].usage.total_tokens == 15 + elif ingress == "responses": + responses_stream: Final = await litellm.aresponses(input="hello", **kwargs) + events: Final = [event async for event in responses_stream] + assert "".join(event.delta for event in events if event.type == "response.output_text.delta") == "hello" + assert events[-1].response.usage.total_tokens == 15 + else: + messages_stream: Final = await litellm.anthropic.messages.acreate( + messages=[{"role": "user", "content": "hello"}], max_tokens=128000, **kwargs + ) + wire_events: Final = b"".join([event async for event in messages_stream]).decode() + message_events: Final = [ + json.loads(line[6:]) for line in wire_events.splitlines() if line.startswith("data: ") + ] + assert ( + "".join( + event["delta"].get("text", "") for event in message_events if event["type"] == "content_block_delta" + ) + == "hello" + ) + assert ( + next(event["usage"]["output_tokens"] for event in message_events if event["type"] == "message_delta") + == 5 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("factory_path", [False, True]) +@pytest.mark.parametrize("async_stream", [False, True]) +async def test_gemini_streams_with_dispatcher_and_factory_configs( + go_catalog: None, factory_path: bool, async_stream: bool +) -> None: + messages: Final = [{"role": "user", "content": "hello"}] + logging_obj: Final = Logging( + model="fixture-gemini", + messages=messages, + stream=True, + call_type="acompletion" if async_stream else "completion", + start_time=datetime(2026, 1, 1), + litellm_call_id="go-stream-regression", + function_id="go-stream-regression", + ) + logging_obj.update_environment_variables( + model="fixture-gemini", + user="", + optional_params={"stream_options": {"include_usage": True}}, + litellm_params={}, + custom_llm_provider="opencode_go", + stream_options={"include_usage": True}, + ) + with respx.mock(using="httpx") as transport: + transport.post( + "https://opencode.ai/inference/go/google/v1beta/models/fixture-gemini:streamGenerateContent?alt=sse" + ).respond( + headers={"content-type": "text/event-stream"}, + text="data: " + json.dumps(upstream_reply("gemini", "fixture-gemini")) + "\n\n", + ) + if factory_path: + result: Final = BaseLLMHTTPHandler().completion( + model="fixture-gemini", + messages=messages, + api_base=None, + custom_llm_provider="opencode_go", + model_response=litellm.ModelResponse(), + encoding=None, + logging_obj=logging_obj, + optional_params={"stream_options": {"include_usage": True}}, + timeout=30, + litellm_params={}, + acompletion=async_stream, + stream=True, + api_key="fake-go-key", + ) + else: + completion: Final = litellm.acompletion if async_stream else litellm.completion + result: Final = completion( + model="opencode_go/fixture-gemini", + messages=messages, + stream=True, + stream_options={"include_usage": True}, + api_key="fake-go-key", + ) + if async_stream: + stream: Final = await result + chunks: Final = [chunk async for chunk in stream] + else: + chunks: Final = tuple(result) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "hello" + assert chunks[-1].usage.total_tokens == 15 + + +@pytest.mark.parametrize("suffix", ["", "/v1beta", "/v1beta/models", "/v1beta/models/"]) +def test_gemini_custom_base_normalizes_version_and_models(go_catalog: None, suffix: str) -> None: + with respx.mock(using="httpx") as transport: + transport.post("https://gateway.example/google/v1beta/models/fixture-gemini:generateContent").respond( + json=upstream_reply("gemini", "fixture-gemini") + ) + response: Final = litellm.completion( + model="opencode_go/fixture-gemini", + api_key="fake-go-key", + api_base=f"https://gateway.example/google{suffix}", + messages=[{"role": "user", "content": "hello"}], + ) + assert response.choices[0].message.content == "hello" + + +def test_gemini_sync_stream_preserves_tool_calls_and_usage(go_catalog: None) -> None: + chunk: Final = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, + "thoughtSignature": "synthetic-signature", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, + } + with respx.mock(using="httpx") as transport: + transport.post( + "https://opencode.ai/inference/go/google/v1beta/models/fixture-gemini:streamGenerateContent?alt=sse" + ).respond(headers={"content-type": "text/event-stream"}, text="data: " + json.dumps(chunk) + "\n\n") + stream: Final = litellm.completion( + model="opencode_go/fixture-gemini", + api_key="fake-go-key", + stream=True, + messages=[{"role": "user", "content": "Weather in Paris?"}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + stream_options={"include_usage": True}, + ) + chunks: Final = tuple(stream) + tool: Final = next( + c.choices[0].delta.tool_calls[0] for c in chunks if c.choices and c.choices[0].delta.tool_calls + ) + assert tool.function.name == "get_weather" + assert json.loads(tool.function.arguments) == {"city": "Paris"} + assert any(c.choices[0].finish_reason == "tool_calls" for c in chunks if c.choices) + assert chunks[-1].usage.total_tokens == 15 + + +GO_REPORTED_USAGE: Final = { + "cost": 0, + "opencode": {"billing_source": "go", "list_cost": 0.00412, "allowance": 0.00412}, +} +ROW_PRICE_FOR_FIXTURE_USAGE: Final = 10 * 0.000002 + 5 * 0.000005 + + +def go_priced_reply(route: str, model: str, usage_extra: Mapping[str, object]) -> Mapping[str, object]: + reply: Final = upstream_reply(route, model) + key: Final = "usageMetadata" if route == "gemini" else "usage" + usage: Final = TypeAdapter(Mapping[str, object]).validate_python(reply[key]) + return {**reply, key: {**usage, **usage_extra}} + + +class CostRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.costs: list[float | None] = [] # mutable-ok: test probe accumulates logged costs + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: # noqa: ANN001 + self.costs.append(kwargs.get("response_cost")) + + +async def logged_cost(recorder: CostRecorder) -> float | None: + for _ in range(50): + if recorder.costs: + return recorder.costs[-1] + await asyncio.sleep(0.05) + raise AssertionError("success callback never ran") + + +@pytest.fixture +def cost_recorder(monkeypatch: pytest.MonkeyPatch) -> CostRecorder: + recorder: Final = CostRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + return recorder + + +@pytest.mark.asyncio +@pytest.mark.parametrize("upstream", ["chat/completions", "messages", "responses", "gemini"]) +@pytest.mark.parametrize("ingress", ["chat", "responses"]) +async def test_reported_list_cost_is_billed_instead_of_rows( + go_catalog: None, cost_recorder: CostRecorder, upstream: str, ingress: str +) -> None: + model: Final = f"fixture-{upstream.split('/')[0]}" + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response(200, json=go_priced_reply(upstream, model, GO_REPORTED_USAGE)) + ) + kwargs: Final = {"model": f"opencode_go/{model}", "api_key": "fake-go-key"} + if ingress == "chat": + response: Final = await litellm.acompletion(messages=[{"role": "user", "content": "hello"}], **kwargs) + else: + response = await litellm.aresponses(input="hello", **kwargs) + assert response._hidden_params["additional_headers"][RESPONSE_COST_HEADER] == pytest.approx(0.00412) + assert await logged_cost(cost_recorder) == pytest.approx(0.00412) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("usage_extra", [{}, {"cost": 0}, {"opencode": {"billing_source": "go"}}]) +async def test_rows_price_the_request_when_no_list_cost_is_reported( + go_catalog: None, cost_recorder: CostRecorder, usage_extra: Mapping[str, object] +) -> None: + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response(200, json=go_priced_reply("chat/completions", "fixture-chat", usage_extra)) + ) + response: Final = await litellm.acompletion( + model="opencode_go/fixture-chat", api_key="fake-go-key", messages=[{"role": "user", "content": "hello"}] + ) + assert RESPONSE_COST_HEADER not in response._hidden_params.get("additional_headers", {}) + assert await logged_cost(cost_recorder) == pytest.approx(ROW_PRICE_FOR_FIXTURE_USAGE) + + +def sse(*events: Mapping[str, object]) -> str: + return "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_chat_stream_bills_the_final_usage_chunks_list_cost( + go_catalog: None, cost_recorder: CostRecorder +) -> None: + chunk: Final = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "fixture-chat", + } + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=sse( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hello"}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + { + **chunk, + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, **GO_REPORTED_USAGE}, + }, + ), + ) + ) + stream: Final = await litellm.acompletion( + model="opencode_go/fixture-chat", + api_key="fake-go-key", + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = [piece async for piece in stream] + assert "".join(piece.choices[0].delta.content or "" for piece in chunks if piece.choices) == "hello" + assert await logged_cost(cost_recorder) == pytest.approx(0.00412) + + +@pytest.mark.asyncio +async def test_responses_stream_bills_the_completed_events_list_cost( + go_catalog: None, cost_recorder: CostRecorder +) -> None: + completed: Final = go_priced_reply("responses", "fixture-responses", GO_REPORTED_USAGE) + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=sse( + { + "type": "response.output_text.delta", + "delta": "hello", + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + }, + {"type": "response.completed", "response": completed, "sequence_number": 2}, + ), + ) + ) + stream: Final = await litellm.aresponses( + model="opencode_go/fixture-responses", api_key="fake-go-key", input="hello", stream=True + ) + events: Final = [event async for event in stream] + assert events[-1].type == "response.completed" + assert events[-1].response._hidden_params["additional_headers"][RESPONSE_COST_HEADER] == pytest.approx(0.00412) + assert await logged_cost(cost_recorder) == pytest.approx(0.00412) + + +def test_stream_chunk_builder_reads_list_cost_from_dict_chunks(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="opencode_go/glm-5.3", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="dict-chunks", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "opencode_go" + chunk = { + "id": "chatcmpl-go", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-5.3", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + } + final_usage = {**chunk, "choices": [], "usage": dict(GO_REPORTED_USAGE)} + + assembled = litellm.stream_chunk_builder(chunks=[chunk, final_usage], logging_obj=logging_obj) + + assert assembled is not None + assert assembled._hidden_params["response_cost"] == GO_REPORTED_USAGE["opencode"]["list_cost"] diff --git a/tests/unit/llms/opencode/test_transformation.py b/tests/unit/llms/opencode/test_transformation.py new file mode 100644 index 00000000000..96203c23122 --- /dev/null +++ b/tests/unit/llms/opencode/test_transformation.py @@ -0,0 +1,90 @@ +from typing import Final + +import pytest + +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header +from litellm.llms.opencode.transformation import OpenCodeProviderAdapter + + +@pytest.mark.parametrize( + "api_base", + [ + "https://opencode.ai/zen/v1", + "https://opencode.ai/zen/go/v1/", + "https://opencode.ai/inference/openai/v1", + "https://opencode.ai/inference/go/openai/v1", + "https://opencode.ai/inference/go/anthropic/v1", + "https://opencode.ai/inference/anthropic/v1", + ], +) +def test_opencode_defaults_affinity_header_to_the_callers_session(api_base: str) -> None: + headers: Final = {"Authorization": "Bearer test-key"} + params: Final = {"metadata": {"session_id": "conversation-123"}} + assert add_provider_affinity_header( + headers, params, OpenCodeProviderAdapter().get_provider_affinity_header(api_base) + ) == { + **headers, + "x-opencode-session": "conversation-123", + } + assert headers == {"Authorization": "Bearer test-key"} + assert params == {"metadata": {"session_id": "conversation-123"}} + + +@pytest.mark.parametrize( + "api_base", + [ + None, + "https://api.openai.com/v1", + "https://opencode.ai.example/zen/go/v1", + "https://opencode.ai@example.com/zen/go/v1", + "https://example.com/opencode.ai/zen/go/v1", + "https://example.com/v1?upstream=https://opencode.ai/zen/go/v1", + "https://opencode.ai/console/v1", + "http://opencode.ai/zen/go/v1", + ], +) +def test_other_destinations_do_not_receive_an_automatic_session_header(api_base: str | None) -> None: + headers: Final = {"Authorization": "Bearer test-key", "x-existing": "unchanged"} + assert ( + add_provider_affinity_header( + headers, + {"litellm_session_id": "conversation-123"}, + OpenCodeProviderAdapter().get_provider_affinity_header(api_base), + ) + == headers + ) + + +@pytest.mark.parametrize( + "params", + [ + {}, + {"litellm_trace_id": "request-123"}, + {"metadata": {"trace_id": "request-123"}}, + { + "litellm_session_id": "request-123", + "litellm_metadata": {"session_id": "request-123", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, + ], +) +def test_opencode_does_not_invent_a_conversation_id(params: dict[str, object]) -> None: + assert ( + add_provider_affinity_header( + {}, params, OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1") + ) + == {} + ) + + +def test_opencode_preserves_explicit_header_and_affinity_configuration() -> None: + assert add_provider_affinity_header( + {"X-OpenCode-Session": "explicit"}, + {"litellm_session_id": "conversation-123"}, + OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1"), + ) == {"X-OpenCode-Session": "explicit"} + assert add_provider_affinity_header( + {}, + {"litellm_session_id": "conversation-123", "provider_affinity_header": "x-custom-session"}, + OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1"), + ) == {"x-custom-session": "conversation-123"} 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 0dba0d7b90f..7cbabb95931 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -19,6 +19,22 @@ def test_full_payload_guardrail_text_excludes_the_server_streaming_marker(): assert json.loads(text) == body, text +def test_full_payload_guardrail_text_excludes_the_injected_litellm_metadata(): + body: Final = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "guardrails": ["g"]} + litellm_metadata: Final = { + "user_api_key_token": "hashed-key", + "user_api_key_user_email": "someone@example.com", + "user_api_key_user_role": "proxy_admin", + } + + text: Final = PassThroughEndpointHandler()._extract_text_for_guardrail( + {**body, "litellm_metadata": litellm_metadata}, + None, + ) + + assert json.loads(text) == body, text + + @pytest.mark.parametrize( "marker", [ diff --git a/tests/unit/llms/runwayml/image_generation/__init__.py b/tests/unit/llms/runwayml/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/runwayml/image_generation/test_runway_image_generation_transformation.py b/tests/unit/llms/runwayml/image_generation/test_runway_image_generation_transformation.py new file mode 100644 index 00000000000..7eec9541b98 --- /dev/null +++ b/tests/unit/llms/runwayml/image_generation/test_runway_image_generation_transformation.py @@ -0,0 +1,87 @@ +from typing import Final +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.runwayml.image_generation.transformation import RunwayMLImageGenerationConfig +from litellm.types.utils import ImageResponse + +TASK_URL = "https://api.dev.runwayml.com/v1/tasks/task_1" + + +@pytest.mark.respx() +def test_output_dict_with_non_string_values_still_yields_the_image(respx_mock): + respx_mock.get(TASK_URL).respond( + json={ + "id": "task_1", + "status": "SUCCEEDED", + "output": [{"url": "https://cdn.example/image.png", "b64_json": None, "width": 1024}], + }, + status_code=200, + ) + created: Final = httpx.Response( + 200, + json={"id": "task_1", "status": "PENDING"}, + request=httpx.Request( + "POST", "https://api.dev.runwayml.com/v1/text_to_image", headers={"Authorization": "Bearer k"} + ), + ) + + image_response: Final = RunwayMLImageGenerationConfig().transform_image_generation_response( + model="gen4_image", + raw_response=created, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert [image.url for image in image_response.data] == ["https://cdn.example/image.png"] + + +def test_non_object_task_body_raises_the_base_error_class(): + created: Final = httpx.Response( + 200, + json=["unexpected"], + request=httpx.Request( + "POST", "https://api.dev.runwayml.com/v1/text_to_image", headers={"Authorization": "Bearer k"} + ), + ) + + with pytest.raises(AttributeError): + RunwayMLImageGenerationConfig().transform_image_generation_response( + model="gen4_image", + raw_response=created, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.asyncio +async def test_async_non_object_task_body_raises_the_base_error_class(): + created: Final = httpx.Response( + 200, + json=["unexpected"], + request=httpx.Request( + "POST", "https://api.dev.runwayml.com/v1/text_to_image", headers={"Authorization": "Bearer k"} + ), + ) + + with pytest.raises(AttributeError): + await RunwayMLImageGenerationConfig().async_transform_image_generation_response( + model="gen4_image", + raw_response=created, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) diff --git a/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py index 12690d5822b..56c0aee3efd 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py @@ -543,8 +543,9 @@ def test_nova_streaming_sends_stream_flag_and_parses_chunks_and_usage(respx_mock "stream": True, "stream_options": {"include_usage": True}, } - assert tuple(chunk.choices[0].delta.content for chunk in chunks) == ("Nova ", "stream", None, None) - assert tuple(chunk.choices[0].finish_reason for chunk in chunks) == (None, None, "stop", None) + assert tuple(chunk.choices[0].delta.content for chunk in chunks[:-1]) == ("Nova ", "stream", None) + assert chunks[-1].choices == [] + assert tuple(chunk.choices[0].finish_reason for chunk in chunks[:-1]) == (None, None, "stop") assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == ( 2, 1, diff --git a/tests/unit/llms/test_api_base_trailing_slash_routes.py b/tests/unit/llms/test_api_base_trailing_slash_routes.py new file mode 100644 index 00000000000..ce8cecf11a5 --- /dev/null +++ b/tests/unit/llms/test_api_base_trailing_slash_routes.py @@ -0,0 +1,119 @@ +import asyncio +import contextlib +from collections.abc import Callable, Iterator, Mapping +from types import MappingProxyType +from typing import Final + +import pytest +import respx + +import litellm +from litellm.evals.main import cancel_run, create_run, get_run + +FAKE_KEY: Final = "sk-fake" +FAKE_AWS: Final[Mapping[str, str]] = MappingProxyType( + {"aws_access_key_id": "AKIAFAKE", "aws_secret_access_key": "fake-secret", "aws_region_name": "us-east-1"} +) +MESSAGES: Final = ({"role": "user", "content": "hi"},) +GATEWAY_PATH: Final = "/gw/" + + +GATEWAY_HOST: Final = "gateway.test" + + +@pytest.fixture +def received_paths() -> Iterator[Callable[[], tuple[str, ...]]]: + with respx.mock(assert_all_called=False) as router: + route: Final = router.route(host=GATEWAY_HOST).respond( + 404, json={"error": {"message": "not found", "code": 404}} + ) + yield lambda: tuple(call.request.url.raw_path.decode().split("?", 1)[0] for call in route.calls) + + +def _drain(result: object) -> None: + if isinstance(result, litellm.CustomStreamWrapper): + for _ in result: + pass + + +def _completion(model: str, stream: bool = False, **kwargs: object) -> Callable[[str], object]: + return lambda api_base: _drain( + litellm.completion(model=model, messages=list(MESSAGES), api_base=api_base, stream=stream, timeout=5, **kwargs) + ) + + +def _embedding(model: str, **kwargs: object) -> Callable[[str], object]: + return lambda api_base: litellm.embedding(model=model, input=["hi"], api_base=api_base, timeout=5, **kwargs) + + +ROUTES: Final[Mapping[str, Callable[[str], object]]] = MappingProxyType( + { + "anthropic-chat": _completion("anthropic/claude-opus-5-5", api_key=FAKE_KEY), + "bedrock-converse": _completion("bedrock/converse/anthropic.claude-sonnet-5-5", **FAKE_AWS), + "bedrock-converse-stream": _completion("bedrock/converse/anthropic.claude-sonnet-5-5", stream=True, **FAKE_AWS), + "gemini-chat": _completion("gemini/gemini-pro-latest", api_key=FAKE_KEY), + "gemini-chat-stream": _completion("gemini/gemini-pro-latest", stream=True, api_key=FAKE_KEY), + "mistral-chat": _completion("mistral/mistral-large-latest", api_key=FAKE_KEY), + "gigachat-chat": _completion("gigachat/GigaChat-2-Max", api_key=FAKE_KEY), + "bedrock-titan-embedding": _embedding("bedrock/amazon.titan-embed-text-v2:0", **FAKE_AWS), + "gemini-embedding": _embedding("gemini/gemini-embedding-2", api_key=FAKE_KEY), + "gigachat-embedding": _embedding("gigachat/Embeddings-2", api_key=FAKE_KEY), + "volcengine-embedding": _embedding("volcengine/doubao-embedding-text-240715", api_key=FAKE_KEY), + "ollama-embedding": _embedding("ollama/nomic-embed-text"), + "oobabooga-chat": _completion("oobabooga/local-model"), + "oobabooga-embedding": _embedding("oobabooga/local-model"), + "bedrock-rerank": lambda api_base: litellm.rerank( + model="bedrock/arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0", + query="hi", + documents=["a"], + api_base=api_base, + **FAKE_AWS, + ), + "bedrock-image-generation": lambda api_base: litellm.image_generation( + model="bedrock/amazon.nova-canvas-v1:0", prompt="a cat", api_base=api_base, timeout=5, **FAKE_AWS + ), + "gemini-interaction-get": lambda api_base: litellm.interactions.get( + interaction_id="interaction-1", api_base=api_base, api_key=FAKE_KEY + ), + "gemini-interaction-delete": lambda api_base: litellm.interactions.delete( + interaction_id="interaction-1", api_base=api_base, api_key=FAKE_KEY + ), + "gemini-interaction-cancel": lambda api_base: litellm.interactions.cancel( + interaction_id="interaction-1", api_base=api_base, api_key=FAKE_KEY + ), + "openai-eval-run-create": lambda api_base: create_run( + eval_id="eval-1", data_source={"type": "jsonl"}, api_base=api_base, api_key=FAKE_KEY + ), + "openai-eval-run-get": lambda api_base: get_run( + eval_id="eval-1", run_id="run-1", api_base=api_base, api_key=FAKE_KEY + ), + "openai-eval-run-cancel": lambda api_base: cancel_run( + eval_id="eval-1", run_id="run-1", api_base=api_base, api_key=FAKE_KEY + ), + "edenai-model-list": lambda api_base: litellm.get_valid_models( + check_provider_endpoint=True, custom_llm_provider="edenai", api_key=FAKE_KEY, api_base=api_base + ), + "gemini-rag-ingest": lambda api_base: asyncio.run( + litellm.aingest( + ingest_options={ + "vector_store": {"custom_llm_provider": "gemini", "api_base": api_base, "api_key": FAKE_KEY} + }, + file_data=("notes.txt", b"hello", "text/plain"), + ) + ), + } +) + + +@pytest.mark.parametrize("route", ROUTES.values(), ids=ROUTES.keys()) +def test_api_base_with_a_trailing_slash_reaches_the_gateway_path_with_single_slashes( + received_paths: Callable[[], tuple[str, ...]], route: Callable[[str], object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("GIGACHAT_ACCESS_TOKEN", FAKE_KEY) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with contextlib.suppress(Exception): + route(f"http://{GATEWAY_HOST}{GATEWAY_PATH}") + + paths: Final = received_paths() + assert paths, "the route never reached the api_base" + assert all(path.startswith(GATEWAY_PATH) and "//" not in path for path in paths), paths diff --git a/tests/unit/llms/test_get_complete_url_trailing_slash.py b/tests/unit/llms/test_get_complete_url_trailing_slash.py new file mode 100644 index 00000000000..9118498c2c2 --- /dev/null +++ b/tests/unit/llms/test_get_complete_url_trailing_slash.py @@ -0,0 +1,278 @@ +import copy +import importlib +import inspect +import itertools +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType, ModuleType +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import pytest +from _pytest.mark.structures import ParameterSet + +import litellm +from litellm.llms.elevenlabs.text_to_speech.transformation import ElevenLabsTextToSpeechConfig +from litellm.llms.meta.realtime.transformation import MUSE_MODEL +from litellm.llms.openai_like.dynamic_config import create_config_class, create_responses_config_class +from litellm.llms.openai_like.json_loader import SimpleProviderConfig + +API_BASE_WITH_TRAILING_SLASH: Final = "https://gw.example.com/p/" +API_BASE_WITHOUT_TRAILING_SLASH: Final = API_BASE_WITH_TRAILING_SLASH.rstrip("/") +LLMS_ROOT: Final = Path(litellm.__file__).parent / "llms" +PACKAGE_ROOT: Final = LLMS_ROOT.parent.parent + +ARGUMENTS_BY_NAME: Final[Mapping[str, object]] = MappingProxyType( + { + "api_key": "sk-test", + "model": "test-model", + "optional_params": {}, + "litellm_params": {"api_key": "sk-test", "vertex_project": "project", "vertex_location": "us-central1"}, + "stream": False, + "endpoint": "v1/messages", + "request_query_params": None, + "data": None, + "api_version": "2024-01-01", + "vector_store_id": "vs_test", + } +) + + +SAMPLE_JSON_PROVIDER: Final = SimpleProviderConfig( + "json-provider", {"base_url": "https://json-provider.example.com", "api_key_env": "UNSET_KEY"} +) +FACTORY_BUILDERS: Final[tuple[type, ...]] = ( + create_config_class(SAMPLE_JSON_PROVIDER), + create_responses_config_class(SAMPLE_JSON_PROVIDER), +) + + +async def _resolve_access_token() -> str: + return "token" + + +@dataclass(frozen=True, slots=True) +class BuilderInputs: + constructor: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + call: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + api_base_option: str | None = None + + +INPUTS_BY_BUILDER: Final[Mapping[str, BuilderInputs]] = MappingProxyType( + { + "anthropic.skills.transformation.AnthropicSkillsConfig": BuilderInputs(call={"endpoint": "skills"}), + "azure_ai.image_edit.transformation.AzureFoundryFluxImageEditConfig": BuilderInputs( + call={"litellm_params": {"api_version": "2025-04-01-preview"}} + ), + "bedrock.chat.agentcore.transformation.AmazonAgentCoreConfig": BuilderInputs( + call={"model": "agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/test-runtime"}, + api_base_option="aws_bedrock_runtime_endpoint", + ), + "bedrock.chat.invoke_agent.transformation.AmazonInvokeAgentConfig": BuilderInputs( + call={"model": "agent/AGENT_ID/ALIAS_ID", "optional_params": {"sessionID": "session-id"}} + ), + "black_forest_labs.image_edit.transformation.BlackForestLabsImageEditConfig": BuilderInputs( + call={"model": "flux-kontext-pro"} + ), + "black_forest_labs.image_generation.transformation.BlackForestLabsImageGenerationConfig": BuilderInputs( + call={"model": "flux-pro-1.1"} + ), + "elevenlabs.text_to_speech.transformation.ElevenLabsTextToSpeechConfig": BuilderInputs( + call={"litellm_params": {ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY: "voice-id"}} + ), + "meta.realtime.transformation.MetaRealtimeConfig": BuilderInputs(call={"model": MUSE_MODEL}), + "openai.evals.transformation.OpenAIEvalsConfig": BuilderInputs(call={"endpoint": "evals"}), + "openai_like.messages.transformation.JSONProviderAnthropicMessagesConfig": BuilderInputs( + constructor={"provider": SAMPLE_JSON_PROVIDER} + ), + "predibase.chat.transformation.PredibaseConfig": BuilderInputs( + call={"litellm_params": {"predibase_tenant_id": "tenant"}} + ), + "ragflow.chat.transformation.RAGFlowConfig": BuilderInputs(call={"model": "ragflow/chat/chat-id/model"}), + "scaledown.chat.transformation.ScaleDownChatConfig": BuilderInputs(call={"model": "classify"}), + "vertex_ai.audio_transcription.realtime_transformation.VertexChirpRealtimeConfig": BuilderInputs( + constructor={"resolve_access_token": _resolve_access_token, "project": "project", "location": None}, + call={"model": "chirp_3"}, + ), + "vertex_ai.interactions.transformation.VertexAIInteractionsConfig": BuilderInputs( + constructor={"mint_access_token": lambda _credentials, _project: ("token", "project")} + ), + "vertex_ai.realtime.transformation.VertexAIRealtimeConfig": BuilderInputs( + constructor={"access_token": "token", "project": "project", "location": "us-central1"} + ), + "vertex_ai.text_to_speech.transformation.VertexAILyriaTextToSpeechConfig": BuilderInputs( + call={"model": "lyria-002"} + ), + } +) + +UNCALLABLE_BUILDERS: Final[Mapping[str, str]] = MappingProxyType( + { + "base_llm.search.transformation.BaseSearchConfig": "raises-NotImplementedError-for-providers-to-override", + "bedrock.image_edit.amazon_nova_canvas_image_edit_transformation.BedrockAmazonNovaCanvasImageEditConfig": ( + "raises-because-BedrockImageEdit-builds-the-url" + ), + "sap.chat.transformation.GenAIHubOrchestrationConfig": "url-comes-from-a-live-SAP-deployment-lookup", + "sap.embed.transformation.GenAIHubEmbeddingConfig": "url-comes-from-a-live-SAP-deployment-lookup", + "triton.completion.transformation.TritonConfig": "api_base-must-be-the-full-generate-or-infer-url", + "vertex_ai.rerank.transformation.VertexAIRerankConfig": "ignores-api_base-and-needs-Vertex-credentials", + } +) + + +API_BASE_IS_THE_FULL_URL: Final[Mapping[str, str]] = MappingProxyType( + { + builder_id: "sends-the-configured-api_base-as-the-request-url" + for builder_id in ( + "base_llm.audio_transcription.transformation.BaseAudioTranscriptionConfig", + "base_llm.chat.transformation.BaseConfig", + "base_llm.completion.transformation.BaseTextCompletionConfig", + "base_llm.embedding.transformation.BaseEmbeddingConfig", + "base_llm.image_generation.transformation.BaseImageGenerationConfig", + "base_llm.image_variations.transformation.BaseImageVariationConfig", + "base_llm.vector_store.transformation.BaseDirectVectorStoreConfig", + "bedrock.realtime.transformation.BedrockRealtimeConfig", + "bedrock.search.transformation.AgentCoreSearchConfig", + "brave.search.transformation.BraveSearchConfig", + "cohere.chat.v2_transformation.CohereV2ChatConfig", + "cohere.embed.transformation.CohereEmbeddingConfig", + "dashscope.image_generation.transformation.DashScopeImageGenerationConfig", + "dataforseo.search.transformation.DataForSEOSearchConfig", + "datarobot.chat.transformation.DataRobotConfig", + "duckduckgo.search.transformation.DuckDuckGoSearchConfig", + "google_pse.search.transformation.GooglePSESearchConfig", + "searchapi.search.transformation.SearchAPIConfig", + "tinyfish.search.transformation.TinyfishSearchConfig", + "vertex_ai.text_to_speech.transformation.VertexAITextToSpeechConfig", + "vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation." + "VertexAIPartnerModelsAnthropicMessagesConfig", + ) + } +) + + +HOST_NAMED_LIKE_A_STRIPPED_SUFFIX_IS_NOT_KEPT: Final[Mapping[str, str]] = MappingProxyType( + { + "deepinfra.rerank.transformation.DeepinfraRerankConfig": "removes-every-openai-substring-on-main-too", + "minimax.videos.transformation.MinimaxVideoConfig": "strips-surface-suffixes-to-the-scheme-on-main-too", + "volcengine.videos.transformation.VolcEngineVideoConfig": "strips-surface-suffixes-to-the-scheme-on-main-too", + "vertex_ai.audio_transcription.realtime_transformation.VertexChirpRealtimeConfig": "returns-a-grpc-host", + "azure.realtime.http_transformation.AzureRealtimeHTTPConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "bedrock.responses.transformation.BedrockOpenAIResponsesConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "bedrock_mantle.chat.claude_transformation.BedrockMantleClaudeChatConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "bedrock_mantle.messages.transformation.BedrockMantleAnthropicMessagesConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "bedrock_mantle.responses.transformation.BedrockMantleResponsesAPIConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "deepseek.messages.transformation.DeepSeekAnthropicMessagesConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "ragflow.chat.transformation.RAGFlowConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "scaledown.chat.transformation.ScaleDownChatConfig": "replaces-the-host-with-a-route-segment-on-main-too", + } +) + + +def _builder_id(cls: type) -> str: + return f"{cls.__module__.removeprefix('litellm.llms.')}.{cls.__qualname__}" + + +def _module_name(path: Path) -> str: + return ".".join(path.relative_to(PACKAGE_ROOT).with_suffix("").parts) + + +def _builder_modules() -> tuple[ModuleType, ...]: + builder_paths: Final = tuple(path for path in LLMS_ROOT.rglob("*.py") if "def get_complete_url" in path.read_text()) + return tuple(importlib.import_module(_module_name(path)) for path in sorted(builder_paths)) + + +def _defined_builders(module: ModuleType) -> tuple[type, ...]: + classes: Final = (*(cls for _, cls in inspect.getmembers(module, inspect.isclass)), *FACTORY_BUILDERS) + return tuple(cls for cls in classes if cls.__module__ == module.__name__ and "get_complete_url" in vars(cls)) + + +def _all_builders() -> tuple[type, ...]: + builders: Final = itertools.chain.from_iterable(_defined_builders(module) for module in _builder_modules()) + return tuple(sorted(builders, key=_builder_id)) + + +def _skip_reason(cls: type) -> str | None: + if getattr(vars(cls)["get_complete_url"], "__isabstractmethod__", False): + return "abstract-get_complete_url" + return UNCALLABLE_BUILDERS.get(_builder_id(cls)) + + +def _builder_param(cls: type) -> ParameterSet: + reason: Final = _skip_reason(cls) + if reason is None: + return pytest.param(cls, id=_builder_id(cls)) + return pytest.param(cls, id=f"{_builder_id(cls)}-skipped-{reason}", marks=pytest.mark.skip(reason=reason)) + + +def _unimplemented(*_args: object, **_kwargs: object) -> object: + raise NotImplementedError + + +def _instantiable(cls: type) -> type: + abstract_methods: Final[frozenset[str]] = getattr(cls, "__abstractmethods__", frozenset()) + if not abstract_methods: + return cls + return type(cls.__name__, (cls,), dict.fromkeys(abstract_methods, _unimplemented)) + + +def _call_arguments(builder: object, inputs: BuilderInputs, api_base: str) -> dict[str, object]: + parameters: Final = inspect.signature(getattr(builder, "get_complete_url")).parameters.values() + named: Final = tuple(p for p in parameters if p.kind not in (p.VAR_POSITIONAL, p.VAR_KEYWORD)) + api_base_options: Final = ( + {} if inputs.api_base_option is None else {"optional_params": {inputs.api_base_option: api_base}} + ) + values: Final = {**ARGUMENTS_BY_NAME, "api_base": api_base, **inputs.call, **api_base_options} + missing: Final = tuple(p.name for p in named if p.default is p.empty and p.name not in values) + assert not missing, f"add a value for {missing} to ARGUMENTS_BY_NAME or INPUTS_BY_BUILDER" + return {p.name: copy.deepcopy(values[p.name]) for p in named if p.name in values} + + +def _built_url(cls: type, api_base: str) -> str: + inputs: Final = INPUTS_BY_BUILDER.get(_builder_id(cls), BuilderInputs()) + builder: Final = _instantiable(cls)(**inputs.constructor) + result: Final = getattr(builder, "get_complete_url")(**_call_arguments(builder, inputs, api_base)) + return str(result[0] if isinstance(result, tuple) else result) + + +def _without_one_trailing_path_slash(url: str) -> str: + parts: Final = urlsplit(url) + return urlunsplit(parts._replace(path=parts.path.removesuffix("/"))) + + +@pytest.mark.parametrize("builder_class", [_builder_param(cls) for cls in _all_builders()]) +def test_get_complete_url_keeps_one_slash_after_a_trailing_slash_api_base(builder_class: type) -> None: + url: Final = _built_url(builder_class, API_BASE_WITH_TRAILING_SLASH) + + assert "//" not in urlsplit(url).path, url + + +@pytest.mark.parametrize("builder_class", [_builder_param(cls) for cls in _all_builders()]) +def test_get_complete_url_sends_the_same_url_with_or_without_a_trailing_slash(builder_class: type) -> None: + with_slash: Final = _built_url(builder_class, API_BASE_WITH_TRAILING_SLASH) + without_slash: Final = _built_url(builder_class, API_BASE_WITHOUT_TRAILING_SLASH) + + if _builder_id(builder_class) in API_BASE_IS_THE_FULL_URL: + assert _without_one_trailing_path_slash(with_slash) == _without_one_trailing_path_slash(without_slash) + else: + assert with_slash == without_slash + + +def _host_preserving_builder_param(cls: type) -> ParameterSet: + reason: Final = HOST_NAMED_LIKE_A_STRIPPED_SUFFIX_IS_NOT_KEPT.get(_builder_id(cls)) + if reason is None: + return _builder_param(cls) + return pytest.param(cls, id=f"{_builder_id(cls)}-skipped-{reason}", marks=pytest.mark.skip(reason=reason)) + + +@pytest.mark.parametrize("api_base", ["https://v1", "https://api/v1", "https://openai/v1/", "https://anthropic/v1"]) +@pytest.mark.parametrize("builder_class", [_host_preserving_builder_param(cls) for cls in _all_builders()]) +def test_get_complete_url_keeps_a_host_named_like_the_suffix_it_strips(builder_class: type, api_base: str) -> None: + url: Final = _built_url(builder_class, api_base) + follows_the_api_base_host: Final = ( + urlsplit(_built_url(builder_class, "https://gateway.test/v1")).netloc == "gateway.test" + ) + + assert urlsplit(url).netloc, url + assert not follows_the_api_base_host or urlsplit(url).netloc == urlsplit(api_base).netloc, url diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py index 20d547c3b63..a02840ed4ef 100644 --- a/tests/unit/llms/test_predibase_transformation.py +++ b/tests/unit/llms/test_predibase_transformation.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import AsyncMock, Mock import httpx @@ -84,6 +85,19 @@ def test_predibase_get_complete_url_stream_and_non_stream(): assert stream_url.endswith("/generate_stream") +@pytest.mark.parametrize(("model", "stream"), [("generate", False), ("generate_stream", True)]) +def test_predibase_get_complete_url_appends_the_route_to_a_model_named_like_it(model: str, stream: bool) -> None: + url: Final = PredibaseConfig().get_complete_url( + api_base="https://serving.example.com", + api_key="test-key", + model=model, + optional_params={"stream": stream}, + litellm_params={"predibase_tenant_id": "tenant-123"}, + ) + + assert url.endswith(f"/llms/{model}/{model}") + + def test_predibase_get_complete_url_missing_tenant_id(): config = PredibaseConfig() diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index f59c584eb90..2a550b4a701 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -278,6 +278,46 @@ class TestGetObjectNameLazyParse: assert exc_info.value.message.startswith("Batch input file line 1 could not be converted") +def _tts_batch_row(**body_extra: object) -> bytes: + body: Final = { + "model": "gemini-2.5-flash-preview-tts", + "messages": [{"role": "user", "content": "Say hi"}], + "audio": {"format": "pcm16", "speech_config": {"voice_name": "Kore", "style": "calm"}}, + **body_extra, + } + return json.dumps({"custom_id": "r-0", "method": "POST", "url": "/v1/chat/completions", "body": body}).encode() + + +def _upload_tts_batch_row(row: bytes, litellm_params: dict[str, object]) -> bytes: + request: Final[CreateFileRequest] = {"file": ("batch.jsonl", row + b"\n", "application/jsonl"), "purpose": "batch"} + return _join_upload_body( + VertexAIFilesConfig().transform_create_file_request( + model="", create_file_data=request, optional_params={}, litellm_params=litellm_params + ) + ) + + +class TestBatchRowUnknownVoiceKeys: + def test_unknown_voice_keys_are_a_400_without_drop_params(self) -> None: + with pytest.raises(VertexAIError, match="voice keys style") as exc_info: + _upload_tts_batch_row(_tts_batch_row(), {}) + assert exc_info.value.status_code == 400 + + @pytest.mark.parametrize( + "row, litellm_params", + [(_tts_batch_row(drop_params=True), {}), (_tts_batch_row(), {"drop_params": True})], + ids=["row", "deployment"], + ) + def test_drop_params_on_the_row_or_deployment_drops_unknown_voice_keys( + self, row: bytes, litellm_params: dict[str, object] + ) -> None: + uploaded: Final = json.loads(_upload_tts_batch_row(row, litellm_params)) + + assert uploaded["request"]["generationConfig"]["speechConfig"] == { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}} + } + + class TestStreamingLaziness: def test_upload_stream_transforms_rows_only_as_they_are_pulled(self) -> None: cfg: Final = VertexAIFilesConfig() diff --git a/tests/unit/llms/vertex_ai/gemini/test_long_decode.py b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py index 4465ad7f3df..3fe9a5c7a92 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_long_decode.py +++ b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py @@ -516,9 +516,16 @@ async def test_unary_answer_cut_at_the_callers_cap_is_returned_with_its_token_in assert response.choices[0].message.provider_specific_fields["continuation_token"] == "dG9rLTE=" +@pytest.mark.parametrize( + "choice_count_request", + [{"max_tokens": 100_000, "n": 2}, {"extra_body": {"generationConfig": {"candidateCount": 2}}}], +) @pytest.mark.parametrize("is_async", [True, False]) @pytest.mark.asyncio -async def test_unary_request_for_more_than_one_choice_is_never_continued(is_async: bool): +async def test_unary_request_for_more_than_one_choice_is_never_continued( + monkeypatch: pytest.MonkeyPatch, is_async: bool, choice_count_request: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "100000") sent: Final[list[dict[str, object]]] = [] # Vertex AI answered 400 "Candidate count greater than 1 is not supported when continuation_token is set" # on 2026-10-13 @@ -526,9 +533,8 @@ async def test_unary_request_for_more_than_one_choice_is_never_continued(is_asyn kwargs: Final = { "model": MODEL, "messages": list(MESSAGES), - "max_tokens": 100_000, - "n": 2, "api_key": "fake-gemini-key", + **choice_count_request, } response: Final = ( @@ -548,6 +554,42 @@ async def test_unary_request_for_more_than_one_choice_is_never_continued(is_asyn ] +@pytest.mark.parametrize( + "one_choice_request", + [ + {"max_tokens": 100_000, "n": 0}, + {"extra_body": {"generationConfig": {"candidateCount": None}}}, + {"extra_body": {"generationConfig": {"candidateCount": 0}}}, + ], +) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_request_whose_candidate_count_google_reads_as_one_is_continued( + monkeypatch: pytest.MonkeyPatch, is_async: bool, one_choice_request: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "100000") + sent: Final[list[dict[str, object]]] = [] + hops: Final = (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2)) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "api_key": "fake-gemini-key", + **one_choice_request, + } + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + if is_async + else litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + ) + + assert sent[1] == {**sent[0], "continuationToken": "dG9rLTE="} + assert response.choices[0].message.content == "hello world" + assert response.choices[0].finish_reason == "stop" + + @pytest.mark.asyncio async def test_unary_failed_continuation_hop_returns_the_answer_so_far_with_its_token(): sent: Final[list[dict[str, object]]] = [] @@ -1295,6 +1337,468 @@ async def test_vertex_long_decode_keeps_the_callers_own_authorization_header_on_ assert seen == ["Bearer mine", "Bearer mine"] +BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" +ARGON: Final = "gemini/gemini-4-argon" +UNMAPPED_GEMINI: Final = "gemini/gemini-99-unreleased" +OUTPUT_LIMIT_KEYS: Final = ("max_output_tokens", "maxOutputTokens") + + +def _output_limits(request: Mapping[str, object]) -> dict[str, object]: + generation_config: Final = request.get("generationConfig", {}) + assert isinstance(generation_config, dict) + return {key: value for key, value in generation_config.items() if key in OUTPUT_LIMIT_KEYS} + + +def _cost_map_max_output_tokens(model: str) -> int: + provider, name = model.split("/", 1) + max_output_tokens: Final = litellm.get_model_info(model=name, custom_llm_provider=provider)["max_output_tokens"] + assert isinstance(max_output_tokens, int) + return max_output_tokens + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_request_with_no_output_limit_keeps_continuing_until_the_env_budget_is_spent( + monkeypatch: pytest.MonkeyPatch, is_async: bool, stream: bool +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + (_sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END), _sse(HOP_2_TRUNCATED), _sse(HOP_2)) + if stream + else ( + httpx.Response(200, json=HOP_1), + httpx.Response(200, json=HOP_2_TRUNCATED), + httpx.Response(200, json=HOP_2), + ) + ) + kwargs: Final = {"model": ARGON, "messages": list(MESSAGES), "stream": stream, "api_key": "fake-gemini-key"} + transport: Final = httpx.MockTransport(_replay(hops, sent)) + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport)) + if is_async + else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport))) + ) + chunks: Final = ([chunk async for chunk in response] if is_async else list(response)) if stream else [] + + assert len(sent) == 2 + assert _output_limits(sent[0]) == {"max_output_tokens": 50} + if stream: + assert _stream_text(chunks) == "hello world" + assert _stream_finish_reasons(chunks) == ["length"] + assert {"continuation_token": "dG9rLTI="} in _delta_fields(chunks) + else: + assert response.choices[0].message.content == "hello world" + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message.provider_specific_fields["continuation_token"] == "dG9rLTI=" + + +@pytest.mark.parametrize( + "budget, model, request_params, expected_limits", + [ + ("50", ARGON, {"max_tokens": 30}, {"max_output_tokens": 30}), + ("50", ARGON, {"max_completion_tokens": 30}, {"max_output_tokens": 30}), + ( + "50", + ARGON, + {"extra_body": {"generationConfig": {"maxOutputTokens": 30}}}, + {"maxOutputTokens": 30}, + ), + (None, ARGON, {}, {}), + ("", ARGON, {}, {}), + ("a lot", ARGON, {}, {}), + ("0", ARGON, {}, {}), + ("-5", ARGON, {}, {}), + ("50", UNMAPPED_GEMINI, {}, {}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_request_that_sets_no_output_limit_of_its_own( + monkeypatch: pytest.MonkeyPatch, + budget: str | None, + model: str, + request_params: Mapping[str, object], + expected_limits: Mapping[str, object], +): + if budget is None: + monkeypatch.delenv(BUDGET_ENV, raising=False) + else: + monkeypatch.setenv(BUDGET_ENV, budget) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + **request_params, + ) + + assert _output_limits(sent[0]) == expected_limits + + +@pytest.mark.parametrize( + "modalities, expected_limits", + [ + (["audio"], {}), + (["text", "audio"], {}), + (["text", "image"], {}), + (["text"], {"max_output_tokens": 50}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_request_asking_for_text_output( + monkeypatch: pytest.MonkeyPatch, modalities: list[str], expected_limits: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + modalities=modalities, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == expected_limits + + +def _ai_studio_model_in_mode(mode: str) -> str: + return next( + name + for name, row in litellm.model_cost.items() + if name.startswith("gemini/") and row.get("mode") == mode and row.get("max_output_tokens") + ) + + +@pytest.mark.parametrize("mode", ["audio_speech", "image_generation"]) +@pytest.mark.asyncio +async def test_env_budget_skips_a_model_the_cost_map_does_not_list_as_a_chat_model( + monkeypatch: pytest.MonkeyPatch, mode: str +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=_ai_studio_model_in_mode(mode), + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {} + + +@pytest.mark.parametrize( + "output_modalities, expected_limits", + [ + ({}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["text"]}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["text", "audio"]}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["audio"]}, {}), + ({"supported_output_modalities": "audio"}, {}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_chat_model_the_cost_map_lists_as_answering_text( + monkeypatch: pytest.MonkeyPatch, output_modalities: Mapping[str, object], expected_limits: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + litellm.register_model( + { + UNMAPPED_GEMINI: { + "litellm_provider": "gemini", + "mode": "chat", + "max_output_tokens": 1000, + **output_modalities, + } + } + ) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=UNMAPPED_GEMINI, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == expected_limits + + +@pytest.mark.parametrize( + "extra_body", + [ + {"generation_config": {"temperature": 0.1}}, + {"generation_config": {"max_output_tokens": 30}}, + {"generationConfig": {"response_modalities": ["AUDIO"]}}, + ], +) +@pytest.mark.asyncio +async def test_env_budget_leaves_a_snake_case_generation_config_from_extra_body_alone( + monkeypatch: pytest.MonkeyPatch, extra_body: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + extra_body=dict(extra_body), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {} + + +def _gemini_chat_models_answering_no_text() -> tuple[str, ...]: + rows: Final = ( + (name, row) + for name, row in litellm.model_cost.items() + if row.get("mode") == "chat" and "text" not in row.get("supported_output_modalities", ("text",)) + ) + return tuple( + name if name.startswith("gemini/") else f"vertex_ai/{name}" + for name, row in rows + if row.get("litellm_provider") in ("gemini", "vertex_ai-language-models") + ) + + +@pytest.mark.asyncio +async def test_env_budget_skips_every_chat_model_the_cost_map_lists_as_answering_no_text( + monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint +): + monkeypatch.setenv(BUDGET_ENV, "50") + models: Final = _gemini_chat_models_answering_no_text() + sent: Final[list[dict[str, object]]] = [] + replies: Final = tuple(httpx.Response(200, json=HOP_2) for _ in models) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model in models: + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + vertex_project="p", + vertex_location="us-central1", + vertex_credentials=token_endpoint.credentials, + client=client, + ) + + assert models + assert {model: _output_limits(request) for model, request in zip(models, sent, strict=True)} == dict.fromkeys( + models, {} + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_env_budget_leaves_a_continuation_token_resume_with_no_limit_alone( + monkeypatch: pytest.MonkeyPatch, stream: bool +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + reply: Final = _sse(HOP_2) if stream else httpx.Response(200, json=HOP_2) + + response: Final = await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + continuation_token="dG9rLTE=", + stream=stream, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((reply,), sent))), + ) + if stream: + _ = [chunk async for chunk in response] + + assert sent[0]["continuationToken"] == "dG9rLTE=" + assert _output_limits(sent[0]) == {} + + +@pytest.mark.parametrize("model", ["gemini/gemini-2.5-flash", "vertex_ai/gemini-2.5-flash", ARGON]) +@pytest.mark.asyncio +async def test_env_budget_above_a_models_output_limit_sends_that_limit_instead( + monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint, model: str +): + budget: Final = 2 * _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(budget)) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + vertex_project="p", + vertex_location="us-central1", + vertex_credentials=token_endpoint.credentials, + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": _cost_map_max_output_tokens(model)} + + +@pytest.mark.asyncio +async def test_env_budget_never_goes_above_the_cost_map_limit_when_a_deployment_overrides_it( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + model_limit: Final = _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(2 * model_limit)) + router: Final = litellm.Router( + model_list=[ + {"model_name": "plain", "litellm_params": {"model": model, "api_key": "fake-gemini-key"}}, + { + "model_name": "raised", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": 4 * model_limit}, + }, + ] + ) + sent: Final[list[dict[str, object]]] = [] + replies: Final = (httpx.Response(200, json=HOP_2), httpx.Response(200, json=HOP_2)) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model_name in ("plain", "raised"): + await router.acompletion(model=model_name, messages=list(MESSAGES), client=client) + + sent_limits: Final = [_output_limits(request).get("max_output_tokens", 0) for request in sent] + assert len(sent_limits) == 2 + assert all(isinstance(limit, int) and limit <= model_limit for limit in sent_limits) + + +@pytest.mark.asyncio +async def test_env_budget_caps_each_deployment_at_its_own_model_info_limit_and_never_at_a_siblings( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + budget: Final = _cost_map_max_output_tokens(model) // 2 + low_limit: Final = budget // 4 + monkeypatch.setenv(BUDGET_ENV, str(budget)) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "low", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": low_limit}, + }, + {"model_name": "plain", "litellm_params": {"model": model, "api_key": "fake-gemini-key"}}, + ] + ) + sent: Final[list[dict[str, object]]] = [] + replies: Final = (httpx.Response(200, json=HOP_2), httpx.Response(200, json=HOP_2)) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model_name in ("plain", "low"): + await router.acompletion(model=model_name, messages=list(MESSAGES), client=client) + + assert [_output_limits(request) for request in sent] == [ + {"max_output_tokens": budget}, + {"max_output_tokens": low_limit}, + ] + + +@pytest.mark.parametrize("deployment_limit", [0, -1]) +@pytest.mark.asyncio +async def test_env_budget_ignores_a_deployment_model_info_limit_that_is_not_positive( + monkeypatch: pytest.MonkeyPatch, deployment_limit: int +): + model: Final = "gemini/gemini-2.5-flash" + budget: Final = _cost_map_max_output_tokens(model) // 2 + monkeypatch.setenv(BUDGET_ENV, str(budget)) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "odd", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": deployment_limit}, + } + ] + ) + sent: Final[list[dict[str, object]]] = [] + + await router.acompletion( + model="odd", + messages=list(MESSAGES), + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": budget} + + +@pytest.mark.asyncio +async def test_env_budget_on_a_direct_call_never_goes_above_the_cost_map_limit_after_a_router_raised_it( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + model_limit: Final = _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(2 * model_limit)) + litellm.Router( + model_list=[ + { + "model_name": "raised", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": 4 * model_limit}, + } + ] + ) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": model_limit} + + +@pytest.mark.asyncio +async def test_env_budget_on_a_direct_call_follows_a_lower_limit_set_with_register_model( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + lowered_limit: Final = _cost_map_max_output_tokens(model) // 4 + monkeypatch.setenv(BUDGET_ENV, str(2 * lowered_limit)) + litellm.register_model({model: {"litellm_provider": "gemini", "max_output_tokens": lowered_limit}}) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": lowered_limit} + + +@pytest.mark.parametrize( + "generation_config", + [None, {"responseModalities": None}, {"responseModalities": {"TEXT": 1}}], +) +@pytest.mark.asyncio +async def test_env_budget_sends_a_generation_config_of_an_unexpected_shape_as_it_came( + monkeypatch: pytest.MonkeyPatch, generation_config: object +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + extra_body={"generationConfig": generation_config}, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert sent[0]["generationConfig"] == generation_config + + def test_initial_prompt_tokens_is_absent_on_a_single_hop_and_the_first_hops_count_otherwise(): assert initial_prompt_tokens_of((HOP_1,)) is None assert initial_prompt_tokens_of((HOP_1, HOP_2)) == 10 diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index b97724319c7..75449eacfd3 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -2,7 +2,7 @@ import base64 from collections.abc import Sequence from itertools import chain from pathlib import Path -from typing import Final +from typing import Final, Literal import httpx import pytest @@ -1650,10 +1650,28 @@ def test_convert_tool_response_with_base64_image(): assert inline_data["data"] == test_image_base64 -def test_gemini_history_nests_multimodal_tool_response_parts(): - """Full history conversion should not emit sibling inline_data tool result parts.""" - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - messages = [ +@pytest.mark.parametrize("provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) +@pytest.mark.parametrize( + "model,nested", + [ + (None, True), + ("gemini-2.5-pro", False), + ("gemini/gemini-2.5-flash", False), + ("vertex_ai/gemini-2.5-flash-lite", False), + ("123456789", False), + ("gemini-3.8-flash", True), + ("gemini/gemini-3.8-flash", True), + ("vertex_ai/gemini-3.8-flash", True), + ("gemini-flash-latest", True), + ], +) +def test_gemini_history_places_tool_media_for_model( + model: str | None, + nested: bool, + provider: Literal["gemini", "vertex_ai", "vertex_ai_beta"], +) -> None: + test_image_base64: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + messages: Final = [ {"role": "user", "content": "Get me an image"}, { "role": "assistant", @@ -1679,24 +1697,53 @@ def test_gemini_history_nests_multimodal_tool_response_parts(): "data": test_image_base64, }, }, + { + "type": "file", + "file": {"file_data": "data:application/pdf;base64,JVBERi0xLjQ="}, + }, ], }, + {"role": "user", "content": "Describe the image and file again"}, ] - contents = gemini_convert_messages_with_history(messages=messages) - - tool_response_parts = contents[-1]["parts"] - assert len(tool_response_parts) == 1 - assert "inline_data" not in tool_response_parts[0] - function_response = tool_response_parts[0]["function_response"] - assert function_response["parts"] == [ + contents: Final = ( + gemini_convert_messages_with_history(messages=messages) + if model is None + else transform_request_body( + messages=messages, + model=model, + custom_llm_provider=provider, + optional_params={}, + litellm_params={}, + cached_content=None, + )["contents"] + ) + media_parts: Final = [ { "inline_data": { "data": test_image_base64, "mime_type": "image/png", } - } + }, + {"inline_data": {"data": "JVBERi0xLjQ=", "mime_type": "application/pdf"}}, ] + function_response: Final = { + "name": "get_image", + "response": {"image_ref": "inline"}, + **({"id": "call_get_image"} if model is not None and nested else {}), + **({"parts": media_parts} if nested else {}), + } + assert contents[-2] == { + "role": "user", + "parts": [ + {"function_response": function_response}, + *([] if nested else media_parts), + ], + } + assert contents[-1] == { + "role": "user", + "parts": [{"text": "Describe the image and file again"}], + } def test_convert_tool_response_text_only(): diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index bb624add673..3bc864cc047 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -16,19 +16,25 @@ from pydantic import BaseModel import litellm from litellm import ModelResponse, completion, embedding, image_generation from litellm._logging import verbose_logger +from litellm.cost_calculator import batch_cost_calculator from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token, parse_prompt_tokens_details from litellm.llms.anthropic.pass_through.messages import handler as anthropic_messages_handler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.common_utils import GEMINI_MAX_ENUM_CHARACTERS, VertexAIError from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + OPENAI_TO_GEMINI_TTS_VOICES, VertexGeminiConfig, _log_gemini_sampling_params_deprecated, + normalize_gemini_speech_config, + reject_unknown_gemini_tts_voice_keys, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata -from litellm.types.utils import ChoiceLogprobs, Usage +from litellm.types.utils import ChoiceLogprobs, ModelInfo, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.integrations.custom_logger import CustomLogger @@ -104,6 +110,242 @@ GEMINI_3_IMAGE_SIZE_MAPPINGS: Final[tuple[tuple[str, str, str], ...]] = ( ) +@pytest.mark.parametrize( + "speech_config", + ( + {"name": "aLlOy"}, + {"speech_config": {"voice_name": "aLlOy"}}, + {"voice_config": {"prebuilt_voice_config": {"voice_name": "aLlOy"}}}, + ), +) +def test_gemini_structured_speech_config_maps_openai_voice_names(speech_config: Mapping[str, object]) -> None: + assert normalize_gemini_speech_config(speech_config) == { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES["alloy"]}} + } + + +@pytest.mark.parametrize( + "speaker", + ( + {"speakerAlias": "Ryan", "speakerId": "alloy"}, + {"speaker_alias": "Ryan", "speaker_id": "alloy"}, + ), +) +def test_gemini_speech_config_accepts_cloud_tts_speaker_entries(speaker: Mapping[str, object]) -> None: + speech_config: Final = { + "multi_speaker_voice_config": { + "speaker_voice_configs": [speaker, {"speaker": "Katie", "voice_config": {"voice_name": "Leda"}}] + } + } + + assert normalize_gemini_speech_config(speech_config) == { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speaker": "Ryan", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES["alloy"]}}, + }, + {"speaker": "Katie", "voiceConfig": {"voiceName": "Leda"}}, + ] + } + } + + +@pytest.mark.parametrize( + ("voice", "expected"), + ( + ( + { + "languageCode": "en-US", + "modelName": "gemini-2.5-flash-tts", + "multiSpeakerVoiceConfig": {"speakerVoiceConfigs": [{"speakerAlias": "Ryan", "speakerId": "Umbriel"}]}, + }, + { + "languageCode": "en-US", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Umbriel"}}} + ] + }, + }, + ), + ( + {"model_name": "gemini-2.5-flash-tts", "voice_config": {"prebuilt_voice_config": {"voice_name": "Leda"}}}, + {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}}, + ), + ({"speech_config": {"modelName": "gemini-2.5-flash-tts", "language_code": "fr-FR"}}, {"languageCode": "fr-FR"}), + ), +) +def test_gemini_speech_config_drops_cloud_tts_model_name( + voice: Mapping[str, object], expected: Mapping[str, object] +) -> None: + assert normalize_gemini_speech_config(voice) == expected + + +def test_gemini_speech_config_drops_non_string_keys_without_adding_voice(): + speech_config: Final = { + "speech_config": { + "language_code": "fr-FR", + "voice_config": {"prebuilt_voice_config": {"voice_name": "Leda", 1: "discard"}}, + }, + "format": "pcm16", + } + + assert normalize_gemini_speech_config(speech_config) == { + "languageCode": "fr-FR", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}, + } + + +def test_gemini_speech_config_keeps_only_speech_config_fields(): + speech_config: Final = { + "voice": "Kore", + "metadata": {"source": "studio"}, + "speech_config": {"language_code": "fr-FR", "style": "calm"}, + } + + assert normalize_gemini_speech_config(speech_config) == { + "languageCode": "fr-FR", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}, + } + + +@pytest.mark.parametrize( + "voice, unknown_keys", + [ + ({"name": "Kore", "style": "calm"}, "style"), + ({"speech_config": {"voice": "Kore", "pitch": 2}}, "pitch"), + ({"speechConfig": {"voiceConfig": {}}, "tone": "x", "pitch": 2}, "pitch, tone"), + ({"speechConfig": {"voiceName": "Kore"}, "speech_config": {"style": "calm"}}, "speech_config"), + ], +) +def test_unknown_gemini_tts_voice_keys_are_rejected_without_drop_params(voice, unknown_keys): + with pytest.raises(litellm.UnsupportedParamsError, match=f"voice keys {unknown_keys}\\."): + reject_unknown_gemini_tts_voice_keys("gemini-3.1-flash-tts-preview", voice, False, "vertex_ai") + + reject_unknown_gemini_tts_voice_keys("gemini-3.1-flash-tts-preview", voice, True, "vertex_ai") + + +def test_chat_audio_speech_config_unknown_keys_are_rejected_unless_dropped(): + audio: Final = {"format": "pcm16", "speech_config": {"voice_name": "Kore", "style": "calm"}} + + with pytest.raises(litellm.UnsupportedParamsError, match="voice keys style\\."): + VertexGeminiConfig().map_openai_params( + non_default_params={"audio": audio}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + + dropped: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"audio": audio}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=True, + ) + assert dropped["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + voice_only: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"audio": {"voice": "Kore", "format": "pcm16", "style": "calm"}}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + assert voice_only["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + +def test_chat_audio_speech_config_unknown_keys_follow_global_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", True) + + dropped: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"audio": {"format": "pcm16", "speech_config": {"voice_name": "Kore", "style": "calm"}}}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + + assert dropped["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + +def _chat_speech_config(speech_config: Mapping[str, object], drop_params: bool) -> object: + return VertexGeminiConfig().map_openai_params( + non_default_params={"audio": {"format": "pcm16", "speech_config": speech_config}}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=drop_params, + )["speechConfig"] + + +def test_chat_audio_drop_params_strips_unknown_fields_inside_speaker_voices(): + speech_config: Final = { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Joe", "voice_config": {"prebuilt_voice_config": {"voice_name": "Puck", "pitch": 2}}}, + {"speaker_alias": "Jane", "speaker_id": "Kore", "style": "calm"}, + ], + "turn_order": "alternate", + } + } + + assert _chat_speech_config(speech_config, drop_params=True) == { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Joe", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}}, + {"speaker": "Jane", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}}, + ] + } + } + + +GEMINI_VOICE_CONFIG: Final = { + "voice": "Kore", + "prebuiltVoiceConfig": {"voiceName": "Kore"}, + "replicatedVoiceConfig": {"mimeType": "audio/wav", "voiceSampleAudio": "UklGRg=="}, +} + + +@pytest.mark.parametrize( + "speech_config", + [ + {"voiceConfig": GEMINI_VOICE_CONFIG, "languageCode": "fr-FR"}, + { + "voice_config": { + "voice": "Kore", + "prebuilt_voice_config": {"voice_name": "Kore"}, + "replicated_voice_config": {"mime_type": "audio/wav", "voice_sample_audio": "UklGRg=="}, + }, + "language_code": "fr-FR", + }, + ], +) +def test_chat_audio_drop_params_keeps_every_voice_config_field_gemini_accepts(speech_config: dict[str, object]): + assert _chat_speech_config(speech_config, drop_params=True) == { + "voiceConfig": GEMINI_VOICE_CONFIG, + "languageCode": "fr-FR", + } + + +def test_chat_audio_reads_snake_case_wrapper_beside_null_camel_case_wrapper(): + mapped: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"audio": {"format": "pcm16", "speechConfig": None, "speech_config": {"voice_name": "Kore"}}}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + + assert mapped["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + +def test_unknown_voice_key_error_stays_bounded_for_large_voice_objects(): + voice: Final = {"name": "Kore", **{f"{index:04d}{'x' * 1000}": index for index in range(2000)}} + + with pytest.raises(litellm.UnsupportedParamsError) as excinfo: + reject_unknown_gemini_tts_voice_keys("gemini-3.1-flash-tts-preview", voice, False, "vertex_ai") + + assert "and 1990 more" in excinfo.value.message + assert len(excinfo.value.message) < 1500 + + def test_top_logprobs(): non_default_params = { "top_logprobs": 2, @@ -1276,6 +1518,7 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): v = VertexGeminiConfig() usage_metadata = { "promptTokenCount": 210, + "cachedContentTokenCount": 54, "candidatesTokenCount": 50, "totalTokenCount": 260, "promptTokensDetails": [ @@ -1297,9 +1540,10 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): usage_metadata = UsageMetadata(**usage_metadata) result = v.calculate_usage(completion_response={"usageMetadata": usage_metadata}) - # prompt details are total - cached per modality - assert result.prompt_tokens_details.text_tokens == 16 # 20 - 4 - assert result.prompt_tokens_details.image_tokens == 140 # (90 + 100) - (40 + 10) + assert result.prompt_tokens_details.text_tokens == 20 + assert result.prompt_tokens_details.image_tokens == 190 + assert result.prompt_tokens_details.cached_tokens == 54 + assert result.prompt_tokens_details.cached_tokens_details == CachedTokensDetails(text_tokens=4, image_tokens=50) # candidates details accumulate duplicate modalities assert result.completion_tokens_details.text_tokens == 20 # 15 + 5 @@ -1307,6 +1551,118 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): assert result.completion_tokens == 50 +def _gemini_usage_with_cached_audio(cached_audio_tokens: int) -> Usage: + cache_tokens_details: Final = ( + ( + {"modality": "TEXT", "tokenCount": 5}, + {"modality": "AUDIO", "tokenCount": cached_audio_tokens}, + ) + if cached_audio_tokens > 0 + else ({"modality": "TEXT", "tokenCount": 5},) + ) + completion_response: Final = { + "usageMetadata": { + "promptTokenCount": 40, + "cachedContentTokenCount": 5 + cached_audio_tokens, + "promptTokensDetails": ( + {"modality": "TEXT", "tokenCount": 20}, + {"modality": "AUDIO", "tokenCount": 20}, + ), + "cacheTokensDetails": cache_tokens_details, + "candidatesTokenCount": 7, + "totalTokenCount": 47, + } + } + return VertexGeminiConfig.calculate_usage(completion_response=completion_response) + + +def test_gemini_raw_usage_metadata_prices_cached_audio_at_its_batch_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_token_batches": 1e-6, + "input_cost_per_audio_token": 6e-6, + "input_cost_per_audio_token_batches": 3e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_token_cost_batches": 4e-7, + "cache_read_input_audio_token_cost": 5e-7, + "cache_read_input_audio_token_cost_batches": 2e-7, + } + usage: Final = _gemini_usage_with_cached_audio(5) + expected_prompt_cost: Final = ( + 15 * model_info["input_cost_per_token_batches"] + + 15 * model_info["input_cost_per_audio_token_batches"] + + 5 * model_info["cache_read_input_token_cost_batches"] + + 5 * model_info["cache_read_input_audio_token_cost_batches"] + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-audio-batch-test-model", + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(expected_prompt_cost) + + +def test_gemini_raw_usage_metadata_prices_cached_audio_at_its_standard_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_audio_token": 6e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_audio_token_cost": 5e-7, + } + usage: Final = _gemini_usage_with_cached_audio(5) + expected_prompt_cost: Final = ( + 15 * model_info["input_cost_per_token"] + + 15 * model_info["input_cost_per_audio_token"] + + 5 * model_info["cache_read_input_token_cost"] + + 5 * model_info["cache_read_input_audio_token_cost"] + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(expected_prompt_cost) + + +def test_gemini_usage_totals_are_subtracted_once_and_audio_cache_rate_matches_audio_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_audio_token": 6e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_audio_token_cost": 6e-6, + } + cached_usage: Final = _gemini_usage_with_cached_audio(5) + uncached_usage: Final = _gemini_usage_with_cached_audio(0) + cached_details: Final = parse_prompt_tokens_details(cached_usage) + cached_prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=cached_usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + uncached_prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=uncached_usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert cached_usage.prompt_tokens_details.text_tokens == 20 + assert cached_usage.prompt_tokens_details.audio_tokens == 20 + assert cached_usage.prompt_tokens_details.cached_tokens_details == CachedTokensDetails( + text_tokens=5, audio_tokens=5 + ) + assert cached_details["text_tokens"] == 15 + assert cached_details["audio_tokens"] == 15 + assert cached_prompt_cost == pytest.approx(uncached_prompt_cost) + + def test_vertex_ai_map_thinking_param_with_budget_tokens_0(): """ If budget_tokens is 0, do not set includeThoughts to True @@ -1894,7 +2250,9 @@ def test_gemini_tool_use(respx_mock: respx.MockRouter) -> None: assert [(call.function.name, json.loads(call.function.arguments)) for call in tool_calls] == [ ("get_weather", {"location": "Lima, Peru"}) ] - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["tool_calls"] + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [ + "tool_calls" + ] @pytest.mark.respx(assert_all_called=True) @@ -3037,6 +3395,14 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped(): async def dispatch_failure_handlers(self, *args, **kwargs): return None + async def dispatch_attempt_failure_handlers( + self, exception: Exception, traceback_exception: str + ) -> None: + return None + + def attempt_scoped_copy(self) -> "DummyLogging": + return self + async def failing_make_call(client=None, **kwargs): raise VertexAIError(status_code=400, message="bad input", headers={}) @@ -5901,7 +6267,7 @@ def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt(): def test_vertex_ai_usage_metadata_with_document_tokens_cached(): - """Test that cached DOCUMENT tokens are correctly subtracted from prompt text tokens.""" + """Test cached DOCUMENT token details and unchanged prompt cost.""" v = VertexGeminiConfig() usage_metadata_dict = { @@ -5925,10 +6291,24 @@ def test_vertex_ai_usage_metadata_with_document_tokens_cached(): completion_response = {"usageMetadata": usage_metadata_dict} result = v.calculate_usage(completion_response=completion_response) - # DOCUMENT cached tokens map to cached_text_tokens, so: - # text_tokens = (8 TEXT + 774 DOCUMENT) - 400 cached = 382 - assert result.prompt_tokens_details.text_tokens == 382, "text_tokens should be (8 + 774) - 400 cached = 382" + assert result.prompt_tokens_details.text_tokens == 782 assert result.prompt_tokens_details.cached_tokens == 400 + assert result.prompt_tokens_details.cached_tokens_details.text_tokens == 400 + + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "cache_read_input_token_cost": 5e-7, + } + prompt_cost, _ = generic_cost_per_token( + model="gemini-document-cache-test-model", + usage=result, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx( + 382 * model_info["input_cost_per_token"] + 400 * model_info["cache_read_input_token_cost"] + ) def test_async_streaming_uses_custom_client(): diff --git a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index 630b2e1eb34..8abd012f9d8 100644 --- a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -5,13 +5,19 @@ Based on the test patterns from other rerank providers and the current Vertex AI import json import os +from datetime import datetime, timedelta +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +import responses +import respx +import litellm from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig from litellm.types.rerank import RerankResponse +from litellm.utils import ProviderConfigManager class TestVertexAIRerankTransform: @@ -646,3 +652,83 @@ class TestVertexAIRerankTransform: second_call = mock_ensure_access_token.call_args_list[1] assert first_call.kwargs["credentials"] == "path/to/credentials.json" assert second_call.kwargs["credentials"] == "path/to/credentials.json" + + +@pytest.fixture +def clear_vertex_rerank_cache(request: pytest.FixtureRequest) -> None: + ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear() + request.addfinalizer( + ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("clear_vertex_rerank_cache") +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_rerank_reuses_scoped_credentials_and_refreshes_expired_tokens( + is_async: bool, +) -> None: + now: Final = datetime(2025, 1, 1) + credentials: Final = tuple( + json.dumps( + { + "type": "authorized_user", + "client_id": identity, + "client_secret": "synthetic-client-secret", + "refresh_token": f"synthetic-refresh-{identity}", + } + ) + for identity in ("first-principal", "second-principal") + ) + with ( + patch("litellm.disable_aiohttp_transport", True), + patch("google.auth._helpers.utcnow", return_value=now) as clock, + responses.RequestsMock(assert_all_requests_are_fired=False) as token_http, + respx.mock as rank_http, + ): + for token in ("initial", "other-credential", "other-project", "renewed"): + token_http.post( + "https://oauth2.googleapis.com/token", + json={"access_token": token, "expires_in": 3600, "token_type": "Bearer"}, + ) + rank_http.post( + url__regex=r"https://discoveryengine\.googleapis\.com/v1/projects/[^/]+/" + r"locations/global/rankingConfigs/default_ranking_config:rank" + ).respond(200, json={"records": [{"id": "0", "score": 0.9}]}) + + async def rank( + credential: str, project: str, expected_token: str, refresh_count: int + ) -> None: + params: Final = { + "model": "vertex_ai/semantic-ranker-default@latest", + "query": f"Query for {expected_token}", + "documents": ["A relevant document"], + "vertex_credentials": credential, + "vertex_project": project, + "caching": False, + "num_retries": 0, + } + result: Final = ( + await litellm.arerank(**params) + if is_async + else litellm.rerank(**params) + ) + assert isinstance(result, RerankResponse) + assert result.results[0]["index"] == 0 + assert len(token_http.calls) == refresh_count + request: Final = rank_http.calls.last.request + assert request.headers["Authorization"] == f"Bearer {expected_token}" + assert request.headers["X-Goog-User-Project"] == project + assert str(request.url) == ( + f"https://discoveryengine.googleapis.com/v1/projects/{project}/" + "locations/global/rankingConfigs/default_ranking_config:rank" + ) + assert json.loads(request.content)["query"] == params["query"] + + await rank(credentials[0], "first-project", "initial", 1) + await rank(credentials[0], "first-project", "initial", 1) + await rank(credentials[1], "first-project", "other-credential", 2) + await rank(credentials[0], "second-project", "other-project", 3) + await rank(credentials[0], "first-project", "initial", 3) + clock.return_value = now + timedelta(hours=2) + await rank(credentials[0], "first-project", "renewed", 4) diff --git a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py index 17a92b31284..fa332cef5d8 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -6,13 +6,17 @@ import os import pathlib import re import traceback -from typing import Final +from datetime import datetime, timedelta +from threading import Event +from typing import Final, Literal from unittest.mock import MagicMock, call, patch import google.auth.exceptions import pytest +import requests +import responses -from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MAX_RECURSE_DEPTH import litellm @@ -977,6 +981,22 @@ class TestVertexBase: assert result_url == f"{gateway_api_base}:embedContent" + def test_check_custom_proxy_vertex_api_base_with_path_and_trailing_slash_appends_endpoint_once(self): + gateway_api_base: Final = "https://gateway.example.com/vertex/v1/projects/p/locations/us-central1/endpoints/123" + + _, result_url = VertexBase()._check_custom_proxy( + api_base=f"{gateway_api_base}/", + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=None, + auth_header="Bearer token123", + url="https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/endpoints/123:predict", + model="123", + ) + + assert result_url == f"{gateway_api_base}:predict" + def test_check_custom_proxy_vertex_api_base_with_version_path_grafts_default_path(self): vertex_base = VertexBase() @@ -2573,3 +2593,120 @@ def test_get_access_token_reraises_malformed_service_account_credentials_without assert [record.getMessage().split(". Error: ")[0] for record in caplog.records] == [ "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information" ] + + +def _authorized_user_cache_credentials(identity: str) -> str: + return json.dumps( + { + "type": "authorized_user", + "client_id": identity, + "client_secret": "synthetic-secret", + "refresh_token": identity, + "quota_project_id": identity, + } + ) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("inferred_project", [False, True], ids=["explicit", "inferred"]) +@pytest.mark.asyncio +async def test_credential_cache_bounds_retention_and_reloads_least_recently_used( + is_async: bool, inferred_project: bool +) -> None: + vertex: Final = VertexBase() + limit: Final = max(0, DEFAULT_MAX_LRU_CACHE_SIZE) + with ( + patch("google.auth._helpers.utcnow", return_value=datetime(2025, 1, 1)), + responses.RequestsMock() as oauth, + ): + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + return 200, {}, json.dumps({"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600}) + + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + + async def access(index: int, infer: bool = False) -> tuple[str, str]: + identity: Final = f"project-{index}" + credentials: Final = _authorized_user_cache_credentials(identity) + project: Final = None if infer else identity + result: Final = ( + await vertex.get_access_token_async(credentials, project) + if is_async + else vertex.get_access_token(credentials, project) + ) + assert result[1] == identity + assert len(vertex._credentials_project_mapping) <= limit + return result + + if limit == 0: + assert await access(0, inferred_project) == ("token-1", "project-0") + assert await access(0, inferred_project) == ("token-2", "project-0") + return + + for index in range(limit): + assert await access(index) == (f"token-{index + 1}", f"project-{index}") + assert await access(0) == ("token-1", "project-0") + assert len(oauth.calls) == limit + assert await access(limit, inferred_project) == (f"token-{limit + 1}", f"project-{limit}") + assert len(oauth.calls) == limit + 1 + if limit > (2 if inferred_project else 1): + assert await access(0) == ("token-1", "project-0") + assert len(oauth.calls) == limit + 1 + evicted: Final = 1 if limit > 1 else 0 + assert await access(evicted) == (f"token-{limit + 2}", f"project-{evicted}") + assert len(oauth.calls) == limit + 2 + + +@pytest.mark.parametrize("mode", ["sync", "async", "background"]) +@pytest.mark.skipif(DEFAULT_MAX_LRU_CACHE_SIZE < 2, reason="Eviction during refresh requires two retained entries") +@pytest.mark.asyncio +async def test_evicted_credentials_finish_refresh_without_being_readmitted( + mode: Literal["sync", "async", "background"], +) -> None: + vertex: Final = VertexBase() + limit: Final = DEFAULT_MAX_LRU_CACHE_SIZE + now: Final = datetime(2025, 1, 1) + refresh_started: Final = Event() + release_refresh: Final = Event() + donor: Final = _authorized_user_cache_credentials("donor") + target: Final = _authorized_user_cache_credentials("target") + with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth: + def token_response(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + body: Final = request.body.decode() if isinstance(request.body, bytes) else request.body + assert isinstance(body, str) + is_target: Final = "refresh_token=target" in body + if is_target and len(oauth.calls) == limit: + refresh_started.set() + assert release_refresh.wait(5), "OAuth refresh was not released" + return 200, {}, json.dumps( + {"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600 if is_target else 86400} + ) + + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + assert vertex.get_access_token(donor, None) == ("token-1", "donor") + assert vertex.get_access_token(donor, None) == ("token-1", "donor") + for index in range(limit - 2): + assert vertex.get_access_token(donor, f"filler-{index}") == (f"token-{index + 2}", f"filler-{index}") + assert vertex.get_access_token(target, "target") == (f"token-{limit}", "target") + clock.return_value = now + timedelta(seconds=3550 if mode == "background" else 7200) + refresh: Final = asyncio.create_task( + asyncio.to_thread(vertex.get_access_token, target, "target") + if mode == "sync" + else vertex.get_access_token_async(target, "target") + ) + try: + assert await asyncio.to_thread(refresh_started.wait, 5), "OAuth refresh did not start" + for index in range(limit - 2): + assert vertex.get_access_token(donor, f"filler-{index}") == (f"token-{index + 2}", f"filler-{index}") + assert vertex.get_access_token(donor, None) == ("token-1", "donor") + assert len(vertex._credentials_project_mapping) == limit + assert (target, "target") not in vertex._credentials_project_mapping + finally: + release_refresh.set() + completed: Final = await refresh + await asyncio.gather(*tuple(vertex._background_refresh_tasks.values())) + expected_token: Final = f"token-{limit if mode == 'background' else limit + 1}" + assert completed == (expected_token, "target") + assert len(oauth.calls) == limit + 1 + assert (target, "target") not in vertex._credentials_project_mapping + assert vertex.get_access_token(target, "target") == (f"token-{limit + 2}", "target") + assert len(vertex._credentials_project_mapping) == limit diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py index dc6f35fcfe5..b3d6d5b952f 100644 --- a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py @@ -1,4 +1,10 @@ import base64 +import io +import math +import struct +import sys +import wave +from pathlib import Path from typing import Final import json from collections.abc import Iterator, Mapping @@ -12,13 +18,61 @@ from pydantic import ValidationError import litellm from litellm.litellm_core_utils.exception_mapping_utils import exception_type +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import OPENAI_TO_GEMINI_TTS_VOICES from litellm.llms.vertex_ai.text_to_speech.transformation import ( VertexAILyriaTextToSpeechConfig, VertexAITextToSpeechConfig, + _fallback_gemini_tts_audio_duration, ) +from litellm.cost_calculator import response_cost_calculator +from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +@pytest.mark.parametrize( + ("voice", "multi_speaker"), + ( + ("aLlOy", False), + ({"name": "aLlOy"}, False), + ({"voice_config": {"prebuilt_voice_config": {"voice_name": "aLlOy"}}}, False), + ( + { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "reader", "voice_config": {"prebuilt_voice_config": {"voice_name": "aLlOy"}}} + ] + } + }, + True, + ), + ), +) +def test_gemini_cloud_tts_maps_openai_voice_names_in_all_voice_forms( + voice: str | dict[str, object], multi_speaker: bool +) -> None: + config: Final = VertexAITextToSpeechConfig() + model: Final = "gemini-3.1-flash-tts-preview" + _, params = config.map_openai_params(model=model, voice=voice, optional_params={"response_format": "mp3"}) + mapped_voice: Final = OPENAI_TO_GEMINI_TTS_VOICES["alloy"] + expected_voice: Final = { + "languageCode": config.DEFAULT_LANGUAGE_CODE, + "modelName": model, + **( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [{"speakerAlias": "reader", "speakerId": mapped_voice}] + } + } + if multi_speaker + else {"name": mapped_voice} + ), + } + + assert params["vertex_voice_dict"] == expected_voice + class TestVertexAITextToSpeechConfig: """Tests for VertexAITextToSpeechConfig transformation""" @@ -48,8 +102,12 @@ class TestVertexAITextToSpeechConfig: assert url == custom_url - @patch.object(VertexAITextToSpeechConfig, "_ensure_access_token") - @patch.object(VertexAITextToSpeechConfig, "_get_token_and_url") + @patch.object( # test-quality-ok: isolates provider credentials while testing request serialization + VertexAITextToSpeechConfig, "_ensure_access_token" + ) + @patch.object( # test-quality-ok: fixes the provider URL at the authentication boundary + VertexAITextToSpeechConfig, "_get_token_and_url" + ) def test_transform_text_to_speech_request_body(self, mock_get_token, mock_ensure_token): """Test that transform_text_to_speech_request generates correct request body""" # Mock authentication @@ -94,6 +152,20 @@ class TestVertexAITextToSpeechConfig: assert "headers" in result assert "Authorization" in result["headers"] + raw_voice_request: Final = config.transform_text_to_speech_request( + model="vertex_ai/chirp", + input="Hello", + voice="en-US-Chirp3-HD-Charon", + optional_params={}, + litellm_params={ + "vertex_credentials": None, + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + headers={}, + ) + assert raw_voice_request["dict_body"]["voice"]["name"] == "en-US-Chirp3-HD-Charon" + def test_voice_mapping_openai_to_vertex(self): """Test that OpenAI voice names are correctly mapped to Vertex AI voices""" config = VertexAITextToSpeechConfig() @@ -131,6 +203,506 @@ class TestVertexAITextToSpeechConfig: assert voice_str is None assert voice_dict == voice_input + def test_gemini_tts_multi_speaker_voice_mapping(self): + config = VertexAITextToSpeechConfig() + + voice = { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + { + "speaker": "Ryan", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Umbriel", + }, + }, + }, + { + "speaker": "Katie", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Leda", + }, + }, + }, + ], + }, + } + + expected_voice = { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speakerAlias": "Ryan", + "speakerId": "Umbriel", + }, + { + "speakerAlias": "Katie", + "speakerId": "Leda", + }, + ], + }, + } + + voice_str, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + assert voice_str is None + assert optional_params["audioEncoding"] == "MP3" + assert optional_params["vertex_voice_dict"] == expected_voice + + voice_str, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "opus"}, + voice=voice, + ) + assert voice_str is None + assert optional_params["audioEncoding"] == "OGG_OPUS" + assert optional_params["vertex_voice_dict"] == expected_voice + + def test_gemini_tts_nested_language_and_voice_alias_reach_cloud_tts(self): + config: Final = VertexAITextToSpeechConfig() + voice: Final = {"speech_config": {"language_code": "fr-FR"}, "voice": "Kore"} + + voice_name, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + + assert voice_name is None + assert optional_params["vertex_voice_dict"] == { + "languageCode": "fr-FR", + "modelName": "gemini-3.1-flash-tts-preview", + "name": "Kore", + } + + @pytest.mark.parametrize( + ("voice", "expected_voice"), + [ + ( + { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Ryan", "voice_config": {"prebuilt_voice_config": {"voice_name": "Umbriel"}}}, + {"speaker": "Katie", "voice_config": {"prebuilt_voice_config": {"voice_name": "Leda"}}}, + ], + }, + "speech_config": {"language_code": "fr-FR"}, + }, + { + "languageCode": "fr-FR", + "modelName": "gemini-3.1-flash-tts-preview", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speakerAlias": "Ryan", "speakerId": "Umbriel"}, + {"speakerAlias": "Katie", "speakerId": "Leda"}, + ], + }, + }, + ), + ( + { + "voice_config": {"prebuilt_voice_config": {"voice_name": "Puck"}}, + "speechConfig": {"languageCode": "en-GB"}, + }, + {"languageCode": "en-GB", "modelName": "gemini-3.1-flash-tts-preview", "name": "Puck"}, + ), + ], + ) + def test_gemini_tts_voice_beside_speech_config_wrapper_reaches_cloud_tts( + self, voice: dict[str, object], expected_voice: dict[str, object] + ): + config: Final = VertexAITextToSpeechConfig() + + voice_name, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + + assert voice_name is None + assert optional_params["vertex_voice_dict"] == expected_voice + + @pytest.mark.parametrize( + ("voice", "expected_language", "expected_voices"), + [ + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Umbriel"}}}, + {"speaker": "Katie", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}}, + ], + }, + "speech_config": { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Ana", "voice_config": {"prebuilt_voice_config": {"voice_name": "Kore"}}}, + {"speaker": "Ben", "voice_config": {"prebuilt_voice_config": {"voice_name": "Puck"}}}, + ], + }, + }, + }, + None, + (("Ana", "Kore"), ("Ben", "Puck")), + ), + ( + {"name": "Puck", "languageCode": "fr-FR", "speech_config": {"language_code": "de-DE"}}, + "de-DE", + ((None, "Puck"),), + ), + ( + {"name": "Puck", "voice_config": {"prebuilt_voice_config": {"voice_name": "Charon"}}}, + None, + ((None, "Charon"),), + ), + ], + ) + def test_gemini_tts_cloud_tts_speaks_the_voice_wav_sends_to_gemini( + self, + voice: dict[str, object], + expected_language: str | None, + expected_voices: tuple[tuple[str | None, str], ...], + ): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import normalize_gemini_speech_config + + _, optional_params = VertexAITextToSpeechConfig().map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + cloud_tts_voice: Final = optional_params["vertex_voice_dict"] + cloud_tts_speakers: Final = cloud_tts_voice.get("multiSpeakerVoiceConfig", {}).get("speakerVoiceConfigs") + cloud_tts_voices: Final = ( + tuple((speaker["speakerAlias"], speaker["speakerId"]) for speaker in cloud_tts_speakers) + if cloud_tts_speakers is not None + else ((None, cloud_tts_voice["name"]),) + ) + wav_config: Final = normalize_gemini_speech_config(voice, drop_unknown_fields=True) + wav_speakers: Final = wav_config.get("multiSpeakerVoiceConfig", {}).get("speakerVoiceConfigs") + wav_voices: Final = ( + tuple( + (speaker["speaker"], speaker["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"]) + for speaker in wav_speakers + ) + if wav_speakers is not None + else ((None, wav_config["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"]),) + ) + + assert cloud_tts_voices == wav_voices == expected_voices + assert cloud_tts_voice["languageCode"] == (expected_language or "en-US") + assert wav_config.get("languageCode") == expected_language + + def test_gemini_tts_string_voice_reaches_cloud_tts(self): + config: Final = VertexAITextToSpeechConfig() + + voice_name, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice="Kore", + ) + + assert voice_name == "Kore" + assert optional_params["vertex_voice_dict"] == { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "name": "Kore", + } + + @pytest.mark.parametrize( + "speaker_voice_configs", + ( + [None, {"speakerAlias": "Katie", "speakerId": "Leda"}], + [{"speakerAlias": "Ryan"}, {"speakerAlias": "Katie", "speakerId": "Leda"}], + [ + {"speaker": "Ryan", "voice_config": {"prebuilt_voice_config": {"name": "Umbriel"}}}, + {"speakerAlias": "Katie", "speakerId": "Leda"}, + ], + "invalid", + ), + ) + def test_gemini_tts_malformed_speaker_entry_is_a_bad_request(self, speaker_voice_configs: object): + config: Final = VertexAITextToSpeechConfig() + + with pytest.raises(litellm.UnsupportedParamsError, match="speaker_voice_configs"): + config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice={"multiSpeakerVoiceConfig": {"speakerVoiceConfigs": speaker_voice_configs}}, + ) + + def test_gemini_tts_malformed_speaker_entry_is_dropped_with_drop_params(self): + config: Final = VertexAITextToSpeechConfig() + + _, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice={ + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speakerAlias": "Ryan", "speakerId": "Umbriel"}, + {"speakerAlias": "Katie"}, + ] + } + }, + drop_params=True, + ) + + assert optional_params["vertex_voice_dict"]["multiSpeakerVoiceConfig"] == { + "speakerVoiceConfigs": [{"speakerAlias": "Ryan", "speakerId": "Umbriel"}] + } + + @pytest.mark.parametrize( + "voice", + [ + {"speechConfig": {"voice": "Kore"}}, + {"speechConfig": {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}}}, + {"speechConfig": {"voiceConfig": {}}}, + {"speech_config": {}}, + {"multi_speaker_voice_config": {"speaker_voice_configs": []}}, + {"languageCode": "en-US"}, + ], + ) + def test_gemini_tts_nested_voice_mapping_sends_only_cloud_tts_voice_fields(self, voice: dict[str, object]): + config: Final = VertexAITextToSpeechConfig() + + _, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + + assert optional_params["vertex_voice_dict"] == { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "name": "Kore", + } + + @pytest.mark.parametrize( + ("voice", "unsendable_field"), + [ + ( + {"voice_config": {"replicated_voice_config": {"mime_type": "audio/wav", "voice_sample_audio": "UklG"}}}, + "voiceConfig.replicatedVoiceConfig", + ), + ( + {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck", "style": "calm"}}}, + "voiceConfig.prebuiltVoiceConfig.style", + ), + ({"speechConfig": {"voiceConfig": {"voice": "Puck"}}}, "voiceConfig.voice"), + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}}, + {"speaker": "Katie", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}, "tone": "x"}, + ] + } + }, + "multiSpeakerVoiceConfig.speakerVoiceConfigs.tone", + ), + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}, "tone": "x"}, + {"speaker": "Katie", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}, "tone": "y"}, + ] + } + }, + "multiSpeakerVoiceConfig.speakerVoiceConfigs.tone", + ), + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speakerAlias": "Ryan", + "speakerId": "Puck", + "voiceConfig": { + "replicatedVoiceConfig": {"mimeType": "audio/wav", "voiceSampleAudio": "AAAA"} + }, + }, + {"speakerAlias": "Katie", "speakerId": "Kore"}, + ] + } + }, + "multiSpeakerVoiceConfig.speakerVoiceConfigs.voiceConfig.replicatedVoiceConfig", + ), + ({"languageCode": ["en-US"], "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}}, "languageCode"), + ({"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": 5}}}, "voiceConfig.prebuiltVoiceConfig.voiceName"), + ({"voiceConfig": "Puck"}, "voiceConfig"), + ({"multiSpeakerVoiceConfig": [{"speaker": "Ryan"}]}, "multiSpeakerVoiceConfig"), + ({"name": 5}, "name"), + ({"speechConfig": {"voice": ["Puck"]}}, "speechConfig.voice"), + ({"speechConfig": "Puck"}, "speechConfig"), + ({"speech_config": [{"voice_config": {"prebuilt_voice_config": {"voice_name": "Puck"}}}]}, "speech_config"), + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speakerAlias": "Ryan", + "speakerId": "Umbriel", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}, + }, + {"speakerAlias": "Katie", "speakerId": "Kore"}, + ] + } + }, + "multiSpeakerVoiceConfig.speakerVoiceConfigs.voiceConfig.prebuiltVoiceConfig", + ), + ], + ) + def test_gemini_tts_cloud_tts_rejects_voice_fields_it_cannot_send( + self, voice: dict[str, object], unsendable_field: str + ): + config: Final = VertexAITextToSpeechConfig() + + with pytest.raises(litellm.UnsupportedParamsError) as exc_info: + config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + + assert f"voice fields {unsendable_field} as mp3" in str(exc_info.value) + + @pytest.mark.parametrize( + ("voice", "voice_name"), + [ + ({"voiceConfig": {"replicatedVoiceConfig": {"mimeType": "audio/wav", "voiceSampleAudio": "UklG"}}}, "Kore"), + ({"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck", "style": "calm"}}}, "Puck"), + ({"voiceConfig": "Puck"}, "Kore"), + ({"speechConfig": "Puck"}, "Kore"), + ], + ) + def test_gemini_tts_cloud_tts_drops_voice_fields_it_cannot_send_with_drop_params( + self, voice: dict[str, object], voice_name: str + ): + config: Final = VertexAITextToSpeechConfig() + + _, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + drop_params=True, + ) + + assert optional_params["vertex_voice_dict"] == { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "name": voice_name, + } + + @pytest.mark.parametrize( + "voice", + [ + {"name": "Umbriel", "modelName": "gemini-2.5-flash-tts"}, + {"name": "Umbriel", "model_name": "gemini-2.5-flash-tts"}, + {"modelName": "gemini-2.5-flash-tts", "model_name": "chirp-3"}, + { + "modelName": "gemini-2.5-flash-tts", + "model_name": "chirp-3", + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + { + "speaker": "Ryan", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Umbriel", + }, + }, + }, + ], + }, + }, + ], + ) + def test_gemini_tts_ignores_voice_model_name_override(self, voice): + config = VertexAITextToSpeechConfig() + routed_model = "gemini-3.1-flash-tts-preview" + + _, optional_params = config.map_openai_params( + model=routed_model, + optional_params={"response_format": "mp3"}, + voice=voice, + ) + assert optional_params["vertex_voice_dict"]["modelName"] == routed_model + + def test_dispatch_maps_gemini_cloud_tts_params_when_provider_config_skipped(self): + config = VertexAITextToSpeechConfig() + handler = MagicMock() + handler.text_to_speech_handler.return_value = "ok" + + result = config.dispatch_text_to_speech( + model="gemini-3.1-flash-tts-preview", + input="Hi", + voice={ + "name": "Umbriel", + "modelName": "gemini-2.5-flash-tts", + }, + optional_params={"response_format": "mp3"}, + litellm_params_dict={}, + logging_obj=MagicMock(), + timeout=10, + extra_headers=None, + base_llm_http_handler=handler, + aspeech=False, + api_base=None, + api_key=None, + ) + + assert result == "ok" + mapped_params = handler.text_to_speech_handler.call_args.kwargs["text_to_speech_optional_params"] + assert mapped_params["audioEncoding"] == "MP3" + assert mapped_params["vertex_voice_dict"]["name"] == "Umbriel" + assert mapped_params["vertex_voice_dict"]["modelName"] == "gemini-3.1-flash-tts-preview" + + @pytest.mark.parametrize( + ("voice", "expected_name"), + [ + ("en-US-Chirp3-HD-Charon", "en-US-Chirp3-HD-Charon"), + ({"name": "en-US-Chirp3-HD-Charon"}, "en-US-Chirp3-HD-Charon"), + (None, None), + ], + ) + def test_dispatch_keeps_pre_mapped_cloud_tts_params( + self, voice: str | dict[str, str] | None, expected_name: str | None + ): + config = VertexAITextToSpeechConfig() + handler = MagicMock() + handler.text_to_speech_handler.return_value = "ok" + optional_params = { + "audioEncoding": "OGG_OPUS", + "vertex_voice_dict": {"languageCode": "en-US", "name": "en-US-Chirp3-HD-Charon"}, + } + + config.dispatch_text_to_speech( + model="chirp", + input="Hi", + voice=voice, + optional_params=optional_params, + litellm_params_dict={}, + logging_obj=MagicMock(), + timeout=10, + extra_headers=None, + base_llm_http_handler=handler, + aspeech=False, + api_base=None, + api_key=None, + ) + + call_kwargs = handler.text_to_speech_handler.call_args.kwargs + assert call_kwargs["voice"] == expected_name + assert call_kwargs["text_to_speech_optional_params"] is optional_params + assert call_kwargs["text_to_speech_optional_params"]["audioEncoding"] == "OGG_OPUS" + @pytest.mark.parametrize( ("audio", "expected_content_type"), @@ -174,6 +746,150 @@ def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled(): assert result.response.content == raw_pcm +@pytest.mark.parametrize("encoding", ["LINEAR16", "PCM", "MP3", "OGG_OPUS", "ALAW", "MULAW"]) +def test_gemini_cloud_tts_response_bills_text_and_audio(encoding: str, monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "gemini-3.1-flash-tts-preview" + input_text: Final = "Hello from Gemini text to speech" + pcm_audio: Final = b"\x00\x00" * 24000 + wav_buffer: Final = io.BytesIO() + with wave.open(wav_buffer, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(24000) + wav_file.writeframes(pcm_audio) + g711_audio: Final = ( + struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", + 36 + 24000, + b"WAVE", + b"fmt ", + 16, + 6 if encoding == "ALAW" else 7, + 1, + 24000, + 24000, + 1, + 8, + b"data", + 24000, + ) + + b"\x00" * 24000 + ) + fixture_dir: Final = Path(__file__).resolve().parents[4] / "audio_tests" + audio_by_encoding: Final = { + "PCM": pcm_audio, + "LINEAR16": wav_buffer.getvalue(), + "MP3": (fixture_dir / "gemini_tts_speech.mp3").read_bytes(), + "OGG_OPUS": (fixture_dir / "gemini_tts_speech.ogg").read_bytes(), + "ALAW": g711_audio, + "MULAW": g711_audio, + } + audio_bytes: Final = audio_by_encoding[encoding] + logger: Final = MagicMock() + sent_request: Final = httpx.Request( + "POST", + _SYNTHESIZE_URL, + json={ + "input": {"text": input_text}, + "audioConfig": {"audioEncoding": encoding, "sampleRateHertz": 24000}, + }, + ) + raw_response: Final = httpx.Response( + 200, json={"audioContent": base64.b64encode(audio_bytes).decode()}, request=sent_request + ) + + result: Final = VertexAITextToSpeechConfig().transform_text_to_speech_response(model, raw_response, logger) + usage: Final = result.usage + assert usage is not None + duration: Final = len(audio_bytes) / 48000 if encoding == "PCM" else calculate_request_duration(audio_bytes) + assert duration is not None + assert _fallback_gemini_tts_audio_duration(audio_bytes, encoding, 24000) == pytest.approx(duration), ( + "Google Cloud AudioEncoding and RFC 7845, checked 2026-09-23: " + "https://cloud.google.com/text-to-speech/docs/reference/rest/v1/AudioEncoding " + "https://www.rfc-editor.org/rfc/rfc7845.html" + ) + assert usage.completion_tokens == math.ceil(duration * 25), ( + "Google Cloud TTS pricing, 2026-09-23: https://cloud.google.com/text-to-speech/pricing" + ) + assert usage.prompt_tokens == litellm.token_counter(model=model, text=input_text) + assert usage.total_tokens == usage.prompt_tokens + usage.completion_tokens + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.audio_tokens == usage.completion_tokens + with monkeypatch.context() as dependencies: + dependencies.setitem(sys.modules, "soundfile", None) + sdk_result: Final = VertexAITextToSpeechConfig().transform_text_to_speech_response(model, raw_response, logger) + assert sdk_result.usage == usage + + model_info: Final = litellm.get_model_info(model, custom_llm_provider="vertex_ai") + cost: Final = response_cost_calculator( + response_object=result, + model=model, + custom_llm_provider="vertex_ai", + call_type="speech", + optional_params={}, + prompt=input_text, + ) + expected_cost: Final = ( + usage.prompt_tokens * model_info["input_cost_per_token"] + + usage.completion_tokens * model_info["output_cost_per_audio_token"] + ) + assert cost == pytest.approx(expected_cost) + assert cost > usage.prompt_tokens * model_info["input_cost_per_token"] + + +@pytest.mark.parametrize("encoding", ["ALAW", "MULAW"]) +@pytest.mark.parametrize("audio", [b"\x12" * 24000, b"RIFF" + b"\x12" * 23996]) +def test_gemini_cloud_tts_returns_raw_g711_audio_when_duration_decoder_is_unavailable(encoding: str, audio: bytes): + logger: Final = MagicMock() + sent_request: Final = httpx.Request( + "POST", + _SYNTHESIZE_URL, + json={ + "input": {"text": "Hello"}, + "audioConfig": {"audioEncoding": encoding, "sampleRateHertz": 24000}, + }, + ) + raw_response: Final = httpx.Response( + 200, json={"audioContent": base64.b64encode(audio).decode()}, request=sent_request + ) + + with patch("litellm.llms.vertex_ai.text_to_speech.transformation.calculate_request_duration", return_value=None): + result: Final = VertexAITextToSpeechConfig().transform_text_to_speech_response( + "gemini-3.1-flash-tts-preview", raw_response, logger + ) + + assert result.response.content == audio + assert result.usage is not None + assert result.usage.prompt_tokens > 0 + assert result.usage.completion_tokens == 25, ( + "Google Cloud TTS pricing, 2026-09-23: https://cloud.google.com/text-to-speech/pricing" + ) + + +def test_gemini_cloud_tts_returns_audio_unbilled_when_duration_is_unmeasurable(): + logger: Final = MagicMock() + sent_request: Final = httpx.Request( + "POST", + _SYNTHESIZE_URL, + json={ + "input": {"text": "Hello"}, + "audioConfig": {"audioEncoding": "MP3", "sampleRateHertz": 24000}, + }, + ) + raw_response: Final = httpx.Response( + 200, json={"audioContent": base64.b64encode(b"invalid").decode()}, request=sent_request + ) + + with patch("litellm.llms.vertex_ai.text_to_speech.transformation.calculate_request_duration", return_value=None): + result: Final = VertexAITextToSpeechConfig().transform_text_to_speech_response( + "gemini-3.1-flash-tts-preview", raw_response, logger + ) + + assert result.response.content == b"invalid" + assert result.usage is None + + class TestVertexAILyriaTextToSpeechConfig: @pytest.mark.parametrize( "model", @@ -929,3 +1645,516 @@ def test_lyria_malformed_item_position_does_not_change_the_mapped_error(payload: custom_llm_provider="vertex_ai", original_exception=exc_info.value, ) + + +_PCM_AUDIO: Final = b"\x00\x00\x01\x00" * 2400 +_GENERATE_CONTENT_RESPONSE: Final = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "inlineData": { + "mimeType": "audio/L16;codec=pcm;rate=24000", + "data": base64.b64encode(_PCM_AUDIO).decode(), + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 10, "totalTokenCount": 15}, +} + + +@pytest.mark.parametrize( + ("model", "response_format"), + ( + ("gemini-3.8-flash-tts", None), + ("gemini-3.1-flash-tts-preview", None), + ("gemini-3.1-flash-tts-preview", "wav"), + ("gemini-3.1-flash-tts-preview", "pcm"), + ), +) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_formats_the_bridge_serves_stay_on_generate_content( + respx_mock: respx.MockRouter, + google_token_endpoint: responses.RequestsMock, + model: str, + response_format: str | None, +) -> None: + route: Final = respx_mock.post( + url__regex=rf"https://aiplatform\.googleapis\.com/.*/models/{model}:generateContent$" + ).mock(return_value=httpx.Response(200, json=_GENERATE_CONTENT_RESPONSE)) + + response: Final = await litellm.aspeech( + model=f"vertex_ai/{model}", + input="Say hello", + voice="Kore", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + **({} if response_format is None else {"response_format": response_format}), + ) + + assert route.call_count == 1 + assert _PCM_AUDIO in response.content + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_reaches_cloud_tts_for_any_gemini_tts_model( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = await litellm.aspeech( + model="vertex_ai/gemini-3.8-flash-tts", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert response.content == audio_bytes + sent: Final = json.loads(route.calls[0].request.content) + assert sent["voice"]["modelName"] == "gemini-3.8-flash-tts" + assert sent["voice"]["name"] == "Kore" + assert sent["audioConfig"]["audioEncoding"] == "MP3" + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_without_voice_keeps_the_gemini_model( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + sent: Final = json.loads(route.calls[0].request.content) + assert sent["voice"] == {"languageCode": "en-US", "modelName": "gemini-3.1-flash-tts-preview", "name": "Kore"} + + +@pytest.mark.parametrize("voice", ("en-US-Studio-O", "en-US-Neural2-F")) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_rejects_a_cloud_tts_voice( + voice: str, respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + with pytest.raises(litellm.BadRequestError, match=f"{voice} is a Cloud Text-to-Speech voice"): + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice=voice, + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_drop_params_swaps_a_cloud_tts_voice_for_a_gemini_voice( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="en-US-Studio-O", + response_format="mp3", + drop_params=True, + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert json.loads(route.calls[0].request.content)["voice"]["name"] == "Kore" + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_without_a_location_uses_the_vertex_default_region( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + monkeypatch.delenv("VERTEX_LOCATION", raising=False) + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post("https://us-texttospeech.googleapis.com/v1/text:synthesize").mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + ) + + assert route.called + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_sends_instructions_as_the_style_prompt( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Have a wonderful day", + voice="Kore", + response_format="mp3", + instructions="Read this slowly and calmly", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + sent: Final = json.loads(route.calls[0].request.content) + assert sent["input"] == {"text": "Have a wonderful day", "prompt": "Read this slowly and calmly"} + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_with_a_numeric_speaking_rate_returns_billed_audio( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + result: Final = await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Hello there", + voice="Kore", + response_format="mp3", + audioConfig={"audioEncoding": "MP3", "speakingRate": 1.25}, + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert json.loads(route.calls[0].request.content)["audioConfig"]["speakingRate"] == 1.25 + assert result.content == audio_bytes + assert result.usage is not None + assert result.usage.completion_tokens > 0 + + +@pytest.mark.parametrize( + ("vertex_location", "synthesize_url"), + ( + ("us-central1", "https://us-texttospeech.googleapis.com/v1/text:synthesize"), + ("us-east5", "https://us-texttospeech.googleapis.com/v1/text:synthesize"), + ("us", "https://us-texttospeech.googleapis.com/v1/text:synthesize"), + ("europe-west4", "https://eu-texttospeech.googleapis.com/v1/text:synthesize"), + ("europe-north1", "https://eu-texttospeech.googleapis.com/v1/text:synthesize"), + ("eu", "https://eu-texttospeech.googleapis.com/v1/text:synthesize"), + ("northamerica-northeast1", "https://northamerica-northeast1-texttospeech.googleapis.com/v1/text:synthesize"), + ("global", _SYNTHESIZE_URL), + ), +) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_keeps_the_deployment_region( + respx_mock: respx.MockRouter, + google_token_endpoint: responses.RequestsMock, + vertex_location: str, + synthesize_url: str, +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(synthesize_url).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location=vertex_location, + ) + + assert route.call_count == 1 + assert response.content == audio_bytes + + +@pytest.mark.parametrize("vertex_location", ["asia-northeast1", "europe-west2", "australia-southeast1"]) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_rejects_a_region_cloud_tts_does_not_serve( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock, vertex_location: str +) -> None: + with pytest.raises(litellm.BadRequestError, match=f"no Gemini TTS location for vertex_location {vertex_location}"): + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location=vertex_location, + ) + + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_router_gemini_tts_mp3_unknown_voice_is_a_bad_request_that_keeps_the_deployment_serving( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + respx_mock.post(_SYNTHESIZE_URL).mock( + side_effect=lambda request: ( + httpx.Response( + 404, json={"error": {"code": 404, "message": "Voice Koree not found.", "status": "NOT_FOUND"}} + ) + if json.loads(request.content)["voice"]["name"] == "Koree" + else httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gemini-tts", + "litellm_params": { + "model": "vertex_ai/gemini-3.1-flash-tts-preview", + "vertex_credentials": _AUTHORIZED_USER, + "vertex_location": "global", + }, + } + ], + num_retries=0, + ) + + with pytest.raises(litellm.BadRequestError, match="Voice Koree not found"): + await router.aspeech(model="gemini-tts", input="Say hello", voice="Koree", response_format="mp3") + response: Final = await router.aspeech(model="gemini-tts", input="Say hello", voice="Kore", response_format="mp3") + + assert response.content == audio_bytes + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_chirp_unknown_voice_stays_not_found( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response( + 404, json={"error": {"code": 404, "message": "Voice en-US-Nope not found.", "status": "NOT_FOUND"}} + ) + ) + + with pytest.raises(litellm.NotFoundError, match="Voice en-US-Nope not found"): + await litellm.aspeech( + model="vertex_ai/chirp", + input="Say hello", + voice={"languageCode": "en-US", "name": "en-US-Nope"}, + vertex_credentials=_AUTHORIZED_USER, + ) + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_rejects_a_deployment_api_base( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + with pytest.raises(litellm.BadRequestError, match="Request wav or pcm"): + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="us-central1", + api_base="https://vertex.gateway.internal", + ) + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_sends_speak_tags_as_text( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + speech_input: Final = "Have a wonderful day" + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input=speech_input, + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + sent: Final = json.loads(route.calls[0].request.content) + assert sent["input"] == {"text": speech_input} + + +@pytest.mark.parametrize("response_format", ("aac", "flac", "pcm16", "WAV")) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_unsupported_format_is_rejected_before_calling_google( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock, response_format: str +) -> None: + with pytest.raises(litellm.BadRequestError, match="Supported response formats: alaw, mp3, mulaw, opus, pcm, wav"): + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format=response_format, + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_sends_both_speakers_to_cloud_tts( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + result: Final = await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Ryan: Hi.\nKatie: Hello.", + voice={ + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Ryan", "voice_config": {"prebuilt_voice_config": {"voice_name": "Umbriel"}}}, + {"speaker": "Katie", "voice_config": {"prebuilt_voice_config": {"voice_name": "Leda"}}}, + ], + }, + }, + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert result.content == audio_bytes + assert result.usage is not None + assert result.usage.prompt_tokens > 0 + assert result.usage.completion_tokens > 0 + assert route.call_count == 1 + sent: Final = json.loads(route.calls[0].request.content) + assert sent["audioConfig"]["audioEncoding"] == "MP3" + assert sent["voice"] == { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speakerAlias": "Ryan", "speakerId": "Umbriel"}, + {"speakerAlias": "Katie", "speakerId": "Leda"}, + ], + }, + } + + +@pytest.mark.parametrize( + ("speech_params", "error"), + ( + ({"voice": ["Kore"]}, "needs voice to be a Gemini voice name"), + ({"voice": ""}, "needs voice to be a Gemini voice name"), + ({"voice": {"name": "Kore", "style": "calm"}}, "does not support the voice keys style"), + ({"voice": {"speechConfig": {"voice": "Kore", "pitch": 2}}}, "does not support the voice keys pitch"), + ({"voice": "Kore", "instructions": ["Read slowly"]}, "needs instructions to be a string"), + ), +) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_malformed_params_are_bad_requests_unless_dropped( + respx_mock: respx.MockRouter, + google_token_endpoint: responses.RequestsMock, + speech_params: dict[str, object], + error: str, +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + request: Final = { + "model": "vertex_ai/gemini-3.1-flash-tts-preview", + "input": "Say hello", + "response_format": "mp3", + "vertex_credentials": _AUTHORIZED_USER, + "vertex_location": "global", + **speech_params, + } + + with pytest.raises(litellm.BadRequestError, match=error): + await litellm.aspeech(**request) + assert not route.called + + response: Final = await litellm.aspeech(**request, drop_params=True) + + sent: Final = json.loads(route.calls[0].request.content) + assert response.content == audio_bytes + assert sent["input"] == {"text": "Say hello"} + assert sent["voice"] == {"languageCode": "en-US", "modelName": "gemini-3.1-flash-tts-preview", "name": "Kore"} + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_beta_lyria_2_serves_the_openai_params_it_ignores( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + route: Final = respx_mock.post(url__regex=r"https://us-central1-aiplatform\.googleapis\.com/.*/lyria-002:predict$").mock( + return_value=httpx.Response(200, json={"predictions": [{"audioContent": "bHlyaWEtMi1hdWRpbw==", "mimeType": "audio/wav"}]}) + ) + + response: Final = await litellm.aspeech( + model="vertex_ai_beta/lyria-002", + input="A bright synth track", + voice="alloy", + response_format="mp3", + speed=1.5, + vertex_credentials=_AUTHORIZED_USER, + vertex_project="test-project", + vertex_location="us-central1", + ) + + assert route.call_count == 1 + assert response.content == b"lyria-2-audio" + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_wav_unknown_voice_keys_are_bad_requests_unless_dropped( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + route: Final = respx_mock.post( + url__regex=r"https://aiplatform\.googleapis\.com/.*/models/gemini-3.1-flash-tts-preview:generateContent$" + ).mock(return_value=httpx.Response(200, json=_GENERATE_CONTENT_RESPONSE)) + request: Final = { + "model": "vertex_ai/gemini-3.1-flash-tts-preview", + "input": "Say hello", + "voice": {"speech_config": {"voice": "Kore", "style": "calm"}, "modelName": "gemini-3.1-flash-tts-preview"}, + "response_format": "wav", + "vertex_credentials": _AUTHORIZED_USER, + "vertex_location": "global", + } + + with pytest.raises(litellm.BadRequestError, match="does not support the voice keys style"): + await litellm.aspeech(**request) + assert not route.called + + response: Final = await litellm.aspeech(**request, drop_params=True) + + assert _PCM_AUDIO in response.content + assert json.loads(route.calls[0].request.content)["generationConfig"]["speechConfig"] == { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}} + } diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index fb1cf4ced2d..0726ebb5ee1 100644 --- a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -632,7 +632,7 @@ class TestVertexGemmaCompletion: assert chunks[0].choices[0].delta.content == "READY" assert chunks[1].choices[0].finish_reason == "stop" if stream_options and stream_options["include_usage"]: - assert chunks[-1].choices[0].delta.content is None + assert chunks[-1].choices == [] assert chunks[-1].usage.total_tokens == 114 assert span.get_metric("_dd.llmobs.total_tokens") == 114 else: diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index f73a8547158..16130364364 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7693,6 +7693,47 @@ class TestMCPDcrBridgeDualCredential: assert "invalid_token" in str(exc_info.value.headers) +@pytest.mark.parametrize("public_path", ["/toolset/analysis/mcp", "/analysis/mcp", "/mcp", "/mcp/", "/mcp/proxy"]) +@pytest.mark.parametrize("root_path", ["", "/litellm"]) +@pytest.mark.parametrize("invalid_token", [False, True]) +def test_gateway_challenge_metadata_contains_the_public_resource( + public_path: str, root_path: str, invalid_token: bool, +) -> None: + from mcp.shared.auth_utils import check_resource_allowed + from starlette.requests import Request + from starlette.types import Scope + + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import _gateway_dcr_challenge + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_aggregate_protected_resource_response, + oauth_protected_resource_root, + ) + + scope: Final[Scope] = { + "type": "http", "method": "POST", "scheme": "https", "path": "/mcp", + "root_path": root_path, "_original_path": f"{root_path}{public_path}", + "headers": [(b"host", b"gateway.example.com")], + } + request: Final = Request(scope) + challenge: Final = _gateway_dcr_challenge(request, "/mcp", None, invalid_token) + header: Final = (challenge.headers or {})["WWW-Authenticate"] + metadata_url: Final = header.split('resource_metadata="', 1)[1].split('"', 1)[0] + toolset: Final = public_path in ("/toolset/analysis/mcp", "/analysis/mcp") + assert metadata_url == ( + f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource" + + ("" if toolset else "/mcp") + ) + metadata: Final = ( + oauth_protected_resource_root(request) if toolset else _build_aggregate_protected_resource_response(request) + ) + assert check_resource_allowed( + requested_resource=f"https://gateway.example.com{root_path}{public_path}", + configured_resource=metadata["resource"], + ) + assert challenge.status_code == 401 + assert ('error="invalid_token"' in header) is invalid_token + + @pytest.mark.asyncio class TestAggregateGatewayDcrChallenge: """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must @@ -10906,7 +10947,19 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc @pytest.mark.asyncio @pytest.mark.parametrize("current_groups", [[], ["replacement-group"]]) async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session_scope(monkeypatch, current_groups): - permission = LiteLLM_ObjectPermissionTable(object_permission_id="current-policy", mcp_servers=["current-server"]) + from litellm.proxy import proxy_server + from litellm.types.mcp_server.mcp_toolset import MCPToolset + + permission = LiteLLM_ObjectPermissionTable( + object_permission_id="current-policy", mcp_toolsets=["session-toolset"] + ) + toolset: Final = MCPToolset( + toolset_id="session-toolset", toolset_name="session-toolset", + tools=[{"server_id": "current-server", "tool_name": "read"}], + ) + tables: Final = SimpleNamespace(litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset]))) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) current = UserAPIKeyAuth( object_permission=permission, object_permission_id="current-policy", @@ -10931,7 +10984,10 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session caller.mcp_toolset_id = "session-toolset" refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True) - assert refreshed.object_permission == permission + assert refreshed.object_permission == permission.model_copy(update={ + "mcp_servers": ["current-server"], "mcp_toolsets": [], + "mcp_tool_permissions": {"current-server": ["read"]}, + }) assert refreshed.object_permission_id == "current-policy" assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == ( "new-team", diff --git a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index f6902c778e5..fdae575276d 100644 --- a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -9,9 +9,10 @@ if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 from exceptiongroup import BaseExceptionGroup import httpx +import httpx2 import pytest from mcp import MCPError -from mcp.types import ErrorData +from mcp.types import CONNECTION_CLOSED, ErrorData from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, @@ -102,6 +103,68 @@ def test_unknown_exception_is_internal(): assert classify_list_exception(ValueError("who knows")).tag == "internal" +@pytest.mark.parametrize("sdk_wrapped", (False, True)) +@pytest.mark.parametrize("nested", (False, True)) +def test_sse_event_limit_is_a_safe_upstream_fault(sdk_wrapped: bool, nested: bool) -> None: + limit: Final = 1024 + decoder_message: Final = f"Server-sent event exceeded the {limit} byte limit." + failure: Final = ( + MCPError(code=CONNECTION_CLOSED, message=f"SSE stream failed: {decoder_message}") + if sdk_wrapped else httpx2.SSEError(decoder_message) + ) + wrapped: Final = BaseExceptionGroup("private-upstream-secret", [failure]) if nested else failure + fault: Final = classify_list_exception(wrapped) + message: Final = ( + f"Upstream MCP SSE event exceeded the {limit} byte per-event limit. " + "Paginate tools/list responses so each event stays below the limit." + ) + assert fault == ServerListFault(tag="upstream_error", sse_event_limit_bytes=limit) + assert fault.public_message == message + assert list_fault_http_status(fault) == 502 + assert outcome_wire_value(fault) == { + "status": "upstream_error", "reason": "sse_event_too_large", "limit_bytes": limit, "message": message, + } + + +@pytest.mark.parametrize( + "failure", + ( + MCPError(code=CONNECTION_CLOSED, message="Connection closed"), + MCPError(code=-32603, message="SSE stream failed: Server-sent event exceeded the 1024 byte limit."), + MCPError(code=CONNECTION_CLOSED, message="SSE stream failed: Server-sent event exceeded the 0 byte limit."), + MCPError(code=CONNECTION_CLOSED, message="SSE stream failed: Server-sent event exceeded the -1 byte limit."), + MCPError(code=CONNECTION_CLOSED, message=( + "SSE stream failed: Server-sent event exceeded the 10000000000000000000 byte limit." + )), + MCPError(code=CONNECTION_CLOSED, message=( + "SSE stream failed: Server-sent event exceeded the 1024 byte limit. private-secret" + )), + httpx2.SSEError("private-secret Server-sent event exceeded the 1024 byte limit."), + httpx2.SSEError("Expected response with content type 'text/event-stream', got 'private-secret'."), + RuntimeError("SSE stream failed: Server-sent event exceeded the 1024 byte limit."), + ), +) +def test_unrelated_errors_do_not_claim_an_sse_event_limit(failure: Exception) -> None: + fault: Final = classify_list_exception(failure) + assert fault.sse_event_limit_bytes is None + assert fault.public_message is None + assert "private-secret" not in str(outcome_wire_value(fault)) + + +@pytest.mark.parametrize("status_code", (401, 403)) +def test_upstream_auth_takes_precedence_over_sse_event_limit(status_code: int) -> None: + request: Final = httpx2.Request("POST", "https://upstream.example/mcp") + response: Final = httpx2.Response(status_code, request=request) + failure: Final = BaseExceptionGroup("upstream errors", [ + httpx2.SSEError("Server-sent event exceeded the 1024 byte limit."), + httpx2.HTTPStatusError("private-secret", request=request, response=response), + ]) + fault: Final = classify_list_exception(failure) + assert fault.tag == ("auth_required" if status_code == 401 else "forbidden") + assert list_fault_http_status(fault) == status_code + assert fault.sse_event_limit_bytes is None + + def test_wire_value_carries_no_prose(): fault = ServerListFault(tag="upstream_error", status_code=500) assert outcome_wire_value(fault) == {"status": "upstream_error", "http_status": 500} diff --git a/tests/unit/proxy/_experimental/mcp_server/test_catalog.py b/tests/unit/proxy/_experimental/mcp_server/test_catalog.py index b383078c0ed..4f6530e5de1 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_catalog.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_catalog.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock @@ -13,6 +13,7 @@ from mcp.types import ( ListResourcesResult, ListResourceTemplatesRequest, ListResourceTemplatesResult, + ListToolsRequest, ListToolsResult, PaginatedRequestParams, Tool, @@ -34,6 +35,104 @@ CatalogKind = Literal["tools", "prompts", "resources", "templates"] OptionalCatalogResult = ListPromptsResult | ListResourcesResult | ListResourceTemplatesResult +@pytest.mark.asyncio +@pytest.mark.parametrize( + "selected,revoked,expected", + [ + ("narrow", False, ("alpha-read_item",)), + ("full", False, ("alpha-read_item", "alpha-write_item", "beta-read_item", "beta-write_item")), + (None, False, ("alpha-read_item", "alpha-write_item", "beta-read_item", "beta-write_item")), + ("narrow", True, ()), + ], +) +async def test_catalog_refresh_keeps_selected_toolset_with_current_grants_and_contents( + monkeypatch: pytest.MonkeyPatch, + selected: str | None, + revoked: bool, + expected: tuple[str, ...], +) -> None: + from fastapi import HTTPException + + from litellm.caching.dual_cache import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.server import apply_toolset_scope + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.types.mcp_server.mcp_toolset import MCPToolset, MCPToolsetTool + + grants: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="catalog-grants", mcp_toolsets=["narrow", "full"] + ) + live_grants: Final = grants.model_copy(update={"mcp_toolsets": ["full"]}) if revoked else grants + tools: Final[tuple[MCPToolsetTool, ...]] = ( + {"server_id": "alpha", "tool_name": "read_item"}, + {"server_id": "alpha", "tool_name": "write_item"}, + {"server_id": "beta", "tool_name": "read_item"}, + {"server_id": "beta", "tool_name": "write_item"}, + ) + toolsets: Final = ( + MCPToolset(toolset_id="narrow", toolset_name="narrow", tools=[tools[0]]), + MCPToolset(toolset_id="full", toolset_name="full", tools=list(tools)), + ) + + async def read_toolsets(where: Mapping[str, Mapping[str, Sequence[str]]]) -> tuple[MCPToolset, ...]: + return tuple(toolset for toolset in toolsets if toolset.toolset_id in where["toolset_id"]["in"]) + + tables: Final = SimpleNamespace( + litellm_objectpermissiontable=SimpleNamespace(find_unique=AsyncMock(return_value=live_grants)), + litellm_mcptoolsettable=SimpleNamespace(find_many=read_toolsets), + ) + database: Final = SimpleNamespace( + db=tables, + writer_db=tables, + get_data=AsyncMock(return_value=UserAPIKeyAuth(object_permission_id=grants.object_permission_id)), + ) + cache: Final = DualCache() + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": []}) + await cache.async_set_cache(key="toolset_perms:narrow", value={"alpha": ["read_item", "write_item"]}) + manager: Final = operations.global_mcp_server_manager + for server_id in ("alpha", "beta"): + monkeypatch.setitem( + manager.registry, + server_id, + MCPServer(server_id=server_id, name=server_id, transport=MCPTransport.http, spec_path="/catalog.yaml"), + ) + + async def tool_handler() -> str: + return "ok" + + for tool in tools: + global_mcp_tool_registry.register_tool( + name=f"{tool['server_id']}-{tool['tool_name']}", + description=tool["tool_name"], + input_schema={"type": "object"}, + handler=tool_handler, + ) + original: Final = UserAPIKeyAuth( + api_key="catalog-key-hash", object_permission=grants, + object_permission_id=grants.object_permission_id, + ) + original.via_virtual_key = True + caller: Final = await apply_toolset_scope(original, selected) if selected is not None else original + before: Final = caller.model_dump() + context: Final = operations.prepare_context(user_api_key_auth=caller) + if revoked: + with pytest.raises(HTTPException, match="does not have access to toolset") as denied: + await catalog.list_gateway_catalog(context, ListToolsRequest(), log_list_tools_to_spendlogs=False) + assert denied.value.status_code == 403 + else: + result: Final = await catalog.list_gateway_catalog( + context, ListToolsRequest(), log_list_tools_to_spendlogs=False + ) + assert isinstance(result, ListToolsResult) + assert tuple(sorted(tool.name for tool in result.tools)) == expected + assert caller.model_dump() == before + assert original.object_permission == grants + + def page(name: str, cursor: str | None = None, revision: str = "stable") -> ListToolsResult: return ListToolsResult( tools=[Tool(name=name, input_schema={"type": "object"})], diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index b0d53f9c93b..34b021df2b9 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -5,6 +5,7 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations import asyncio import logging import sys +from typing import Final from unittest.mock import AsyncMock, MagicMock import httpx @@ -531,12 +532,12 @@ async def test_scope_challenge_replaces_pending_stream_headers_once(): send = AsyncMock() response = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp") - await response({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(response({"type": "http.response.start", "status": 200, "headers": []})) body = asyncio.create_task(response({"type": "http.response.body", "body": b"event: message\n"})) await asyncio.sleep(0) send.assert_not_awaited() response.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) - await asyncio.wait_for(body, 1) + await asyncio.wait_for(asyncio.gather(headers, body), 1) await response({"type": "http.response.body", "body": b"private tool error"}) assert send.await_count == 2 start, result = [call.args[0] for call in send.await_args_list] @@ -567,7 +568,7 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=request) token = active_mcp_request_ctx_var.set(context) try: - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) await asyncio.sleep(0) for method, status in [("initialize", 200), ("tools/list", 200), ("tools/call", 403)]: @@ -582,7 +583,7 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re upstream = MCPClient(server_url="https://upstream.example/mcp") async with upstream._create_httpx_client_factory(transport=httpx2.MockTransport(lambda request: httpx2.Response(200, content=b"tool result")))() as client: async with client.stream("POST", "https://upstream.example/mcp", json={"method": "tools/call"}): - await asyncio.wait_for(progress, 1) + await asyncio.wait_for(asyncio.gather(headers, progress), 1) assert send.await_args_list[-1].args[0]["body"] == b"progress" gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) await gate({"type": "http.response.body", "body": b"result"}) @@ -594,21 +595,105 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re @pytest.mark.asyncio @pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) -async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation(failure): +@pytest.mark.parametrize("fail_send", [False, True]) +async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation( + failure: type[BaseException], fail_send: bool, +) -> None: from mcp.server.context import ServerRequestContext from starlette.requests import Request from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response - send = AsyncMock() - gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") - context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + send: Final = AsyncMock(side_effect=BrokenPipeError("peer disconnected") if fail_send else None) + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + context: Final = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + await asyncio.sleep(0) with pytest.raises(failure): await finish_scope_response(context, AsyncMock(side_effect=failure)) - await gate({"type": "http.response.start", "status": 200, "headers": []}) + if fail_send: + with pytest.raises(BrokenPipeError, match="peer disconnected"): + await headers + else: + await headers await gate({"type": "http.response.body", "body": b"error"}) assert not gate.pending - assert send.await_args_list[-1].args[0]["body"] == b"error" + assert send.await_args_list[-1].args[0]["type"] == ("http.response.start" if fail_send else "http.response.body") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_send", [False, True]) +async def test_scope_transport_owns_header_send_cancellation_and_errors(fail_send: bool) -> None: + import httpx2 + from starlette.types import Message + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import record_upstream_tool_authorization + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send_started: Final = asyncio.Event() + send_stopped: Final = asyncio.Event() + writable: Final = asyncio.Event() + sent: Final[list[Message]] = [] + + async def send(message: Message) -> None: + sent.append(message) + assert message["type"] == "http.response.start" + send_started.set() + try: + if fail_send: + raise BrokenPipeError("peer disconnected") + await writable.wait() + finally: + send_stopped.set() + + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + body: Final = asyncio.create_task(gate({"type": "http.response.body", "body": b"result"})) + await asyncio.sleep(0) + assert not send_started.is_set() + response: Final = httpx2.Response(200, request=httpx2.Request( + "POST", "https://upstream.example/mcp", json={"method": "tools/call"}, + )) + try: + await record_upstream_tool_authorization(response, pending=gate) + await send_started.wait() + if fail_send: + with pytest.raises(BrokenPipeError, match="peer disconnected"): + await headers + else: + headers.cancel() + with pytest.raises(asyncio.CancelledError): + await headers + await body + assert len(sent) == 1 + assert send_stopped.is_set() + finally: + writable.set() + await send_stopped.wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("allowed_before_start", [False, True]) +async def test_scope_authorization_sends_stream_headers_before_the_first_body_chunk(allowed_before_start: bool) -> None: + from starlette.types import Message + + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send: Final = AsyncMock() + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + start: Final[Message] = {"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")]} + if allowed_before_start: + gate.allow() + headers: Final = asyncio.create_task(gate(start)) + await asyncio.sleep(0) + if not allowed_before_start: + send.assert_not_awaited() + gate.allow() + await headers + send.assert_awaited_once_with(start) + await gate({"type": "http.response.body", "body": b"event: message\n", "more_body": True}) + assert send.await_args_list[-1].args[0]["body"] == b"event: message\n" + assert send.await_count == 2 @pytest.mark.asyncio @@ -697,13 +782,14 @@ async def test_scope_challenges_are_isolated_between_concurrent_requests(): context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) token = active_mcp_request_ctx_var.set(context) try: - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) await asyncio.sleep(0) assert get_scope_response() is gate if denied: gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) else: await record_upstream_tool_authorization(httpx2.Response(200, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": "tools/call"})), pending=get_scope_response()) + await headers await gate({"type": "http.response.body", "body": b"done"}) return send.await_args_list[0].args[0]["status"] finally: 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 7315315a619..ecd08e17cda 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 @@ -7681,7 +7681,7 @@ class TestMCPServerManager: request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}), ) token = active_mcp_request_ctx_var.set(context) - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) async def upstream_call(*args, **kwargs): @@ -7701,7 +7701,7 @@ class TestMCPServerManager: assert not result.is_error finally: gate.allow() - await asyncio.wait_for(progress, 1) + await asyncio.wait_for(asyncio.gather(headers, progress), 1) active_mcp_request_ctx_var.reset(token) @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index b8e28c3f9b8..a873e2ab1ae 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -280,7 +280,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror( send = AsyncMock() response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/pt") request = Request({"type": "http", "headers": [], SCOPE_RESPONSE_KEY: response}) - await response({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(response({"type": "http.response.start", "status": 200, "headers": []})) try: from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.server import ( @@ -328,6 +328,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror( info_calls = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args] assert any("Upstream auth failure" in m for m in info_calls) response.allow() + await headers await response({"type": "http.response.body", "body": b"tool error"}) start = send.await_args_list[0].args[0] assert start["status"] == (403 if required_scope is not None else 200) diff --git a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py index fb676ae9a61..2b682fda982 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py @@ -588,6 +588,6 @@ def test_registered_inference_routes_have_an_explicit_managed_access_decision(ro "/vector_stores", "/vector_store/", "/search", "/containers", "/skills", "/claude-code/", "/interactions", "/agents", "/responses/{", "/responses/input_tokens", "/realtime/client_secrets", "/realtime/calls", "/realtime/transcription_sessions", - )) or normalized in ("/models", "/cursor/models", "/cursor/v1/models") + )) or normalized in ("/models", "/cursor/models", "/cursor/v1/models", "/claude_code_gateway/v1/models") concrete: Final = route.split("?")[0].replace("{model}", "model").replace("{model_name:path}", "model") assert managed_agent_route_allowed(concrete, None) is not unsupported, route diff --git a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py index 8ac0e2a6f6d..8c2d8e4f157 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -813,3 +813,95 @@ def test_messages_gated_by_enable_flag(): with _gateway_env(enabled=False) as (client, _): resp = client.post("/claude_code_gateway/v1/messages", json={"model": "claude-sonnet-4-5", "messages": []}) assert resp.status_code == 404 + + +@pytest.mark.parametrize("headers", [{}, {"anthropic-version": "2023-06-01"}]) +def test_gateway_model_discovery_preserves_grants_and_format( + monkeypatch: pytest.MonkeyPatch, headers: dict[str, str] +) -> None: + from litellm import Router + from litellm.proxy import proxy_server + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + model_router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": f"openai/{name}", "api_key": "local-placeholder"}, + } + for name in ("discovery-example", "restricted-example") + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", model_router) + monkeypatch.setattr(proxy_server, "llm_model_list", model_router.model_list) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "general_settings", {"enable_claude_code_gateway": True}) + caller: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + models=["discovery-example"], + ) + app: Final = FastAPI() + app.include_router(proxy_server.router) + app.include_router(gateway_endpoints.router) + app.dependency_overrides[gateway_endpoints.user_api_key_auth] = lambda: caller + with TestClient(app) as client: + standard: Final = client.get("/v1/models", headers=headers) + gateway: Final = client.get("/claude_code_gateway/v1/models", headers=headers) + assert standard.status_code == 200, standard.text + assert gateway.status_code == 200, gateway.text + assert gateway.json() == standard.json() + assert [model["id"] for model in gateway.json()["data"]] == ["discovery-example"] + invalid_scope: Final = client.get("/claude_code_gateway/v1/models?scope=invalid", headers=headers) + assert invalid_scope.status_code == 400 + + +def test_gateway_model_discovery_requires_authentication() -> None: + with _gateway_env(real_auth=True) as (client, _), pytest.raises(ProxyException) as exc_info: + client.get("/claude_code_gateway/v1/models") + assert exc_info.value.code == "401" + + +@pytest.mark.parametrize("settings", [{}, {"enable_claude_code_gateway": False}]) +def test_gateway_model_discovery_disabled(monkeypatch: pytest.MonkeyPatch, settings: dict[str, bool]) -> None: + from litellm.proxy import proxy_server + + with _gateway_env(real_auth=True) as (client, _): + monkeypatch.setattr(proxy_server, "general_settings", settings) + response: Final = client.get( + "/claude_code_gateway/v1/models", headers={"Authorization": f"Bearer {_MASTER_KEY}"} + ) + assert response.status_code == 404, response.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/models", "/v1/models", "/claude_code_gateway/v1/models"]) +async def test_gateway_model_discovery_preserves_route_exemptions(monkeypatch: pytest.MonkeyPatch, route: str) -> None: + from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + + import litellm + from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.auth.route_checks import RouteChecks + + monkeypatch.setattr(litellm, "max_budget", 1.0) + result: Final = await common_checks( + request_body={}, + team_object=LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0), + user_object=None, + end_user_object=None, + global_proxy_spend=2.0, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=UserAPIKeyAuth(team_id="test-team"), + request=MagicMock(), + ) + assert result is True + assert RouteChecks.is_unbilled_route(route) + assert RouteChecks.is_llm_api_route(route) + monkeypatch.setenv("DISABLE_LLM_API_ENDPOINTS", "true") + monkeypatch.setenv("DISABLE_ADMIN_ENDPOINTS", "true") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + EnterpriseRouteChecks.should_call_route(route) diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index d83a8aa2fcd..e21491fd996 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -1,5 +1,7 @@ import asyncio import json +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -29,6 +31,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER from litellm.exceptions import BudgetExceededError +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( ModelAccessDeniedProxyException, ProxyErrorTypes, @@ -812,6 +815,65 @@ def _http_request(client_host: str | None = "10.1.2.3", headers: dict[str, str] ) +class _AuthFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.data: dict[str, object] | None = None + self.error: Exception | None = None + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> HTTPException | None: + self.data = dict(request_data) + self.error = original_exception + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_data", [[], None, "metadata", 1, True]) +@pytest.mark.parametrize("error_kind", ["auth", "budget", "database"]) +async def test_auth_error_preserves_classification_with_non_object_context( + monkeypatch: pytest.MonkeyPatch, request_data: object, error_kind: str +) -> None: + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(litellm, "budget_exceeded_status_code", 429) + errors: Final = { + "auth": HTTPException(status_code=401, detail="Invalid API key"), + "budget": BudgetExceededError(current_cost=2, max_budget=1), + "database": httpx.ConnectError("Database unavailable"), + } + error: Final = errors[error_kind] + expected: Final = _as_proxy_exception(error) + recorder: Final = _AuthFailureRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(proxy_server, "general_settings", {"allow_requests_on_db_unavailable": False}) + + with pytest.raises(ProxyException) as exc_info: + await UserAPIKeyAuthExceptionHandler.handle_authentication_error( + error, + _http_request(headers={"user-agent": "shape-regression"}), + request_data, + "/key/info", + None, + "sk-test", + ) + + assert exc_info.value.code == expected.code + assert exc_info.value.type == expected.type + assert exc_info.value.message == expected.message + assert recorder.error is error + assert recorder.data is not None + metadata: Final = recorder.data["metadata"] + assert isinstance(metadata, Mapping) + assert metadata["requester_ip_address"] == "10.1.2.3" + assert metadata["user_agent"] == "shape-regression" + + @pytest.mark.asyncio @pytest.mark.parametrize( "auth_error, general_settings, request_kwargs, expected_ip", diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index ee1ae43fd3a..10a0f974a42 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3626,6 +3626,18 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: ) +@pytest.mark.parametrize("allow_client_side_credentials", [False, True]) +@pytest.mark.parametrize("nested", [False, True]) +def test_request_cannot_select_chatgpt_auth_file(allow_client_side_credentials: bool, nested: bool) -> None: + selector: Final = {"chatgpt_auth_file": "/tokens/account-b/auth.json"} + with pytest.raises(ValueError, match="chatgpt_auth_file"): + is_request_body_safe( + request_body={"model": "chatgpt-account-a", **({"extra_body": selector} if nested else selector)}, + general_settings={"allow_client_side_credentials": allow_client_side_credentials}, + llm_router=None, + model="chatgpt-account-a", + ) + class TestIsRequestBodySafeBlocksBedrockProjectOverride: """``aws_bedrock_project_id`` pins a deployment to a Bedrock project so that project's data-retention policy applies to its requests. A @@ -5304,3 +5316,27 @@ async def test_a_non_numeric_max_request_size_mb_is_ignored_without_premium(): patch("litellm.proxy.proxy_server.premium_user", False), ): assert await check_if_request_size_is_safe(request=_request_with_content_length(1024)) is True + + +@pytest.mark.parametrize("allow_client_side_credentials", [False, True]) +@pytest.mark.parametrize("route", ["/model/new", "/model/update", "/model/account-a/update", "/health/test_connection"]) +def test_chatgpt_model_configuration_reaches_role_authorization( + allow_client_side_credentials: bool, route: str +) -> None: + body: Final = {"litellm_params": {"model": "chatgpt/gpt-6-luna", "chatgpt_auth_file": "/tokens/account-a/auth.json"}} + assert is_request_body_safe( + request_body=body, + general_settings={"allow_client_side_credentials": allow_client_side_credentials}, + llm_router=None, + model="", + route=route, + ) + for inference_route in ("/v1/chat/completions", "/v1/responses"): + with pytest.raises(ValueError, match="chatgpt_auth_file"): + is_request_body_safe( + request_body=body, + general_settings={"allow_client_side_credentials": allow_client_side_credentials}, + llm_router=None, + model="", + route=inference_route, + ) diff --git a/tests/unit/proxy/auth/test_mcp_ip_filtering.py b/tests/unit/proxy/auth/test_mcp_ip_filtering.py index 4c219760762..c403daf733f 100644 --- a/tests/unit/proxy/auth/test_mcp_ip_filtering.py +++ b/tests/unit/proxy/auth/test_mcp_ip_filtering.py @@ -6,6 +6,7 @@ external callers only see servers with available_on_public_internet=True. """ import logging +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -176,6 +177,19 @@ class TestMCPClientIPExtraction: assert result == "203.0.113.5" + def test_non_boolean_use_x_forwarded_for_never_trusts_xff(self): + request: Final = _make_request("10.0.0.5", "203.0.113.9, 10.0.0.5") + + result: Final = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": "false", + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + + assert result == "10.0.0.5" + def _make_request(client_host, xff): request = MagicMock(spec=Request) diff --git a/tests/unit/proxy/auth/test_password_policy.py b/tests/unit/proxy/auth/test_password_policy.py index f9f5025b57f..0ad1083ebe0 100644 --- a/tests/unit/proxy/auth/test_password_policy.py +++ b/tests/unit/proxy/auth/test_password_policy.py @@ -180,6 +180,18 @@ async def test_breach_check_skipped_when_disabled(): assert result is None +@pytest.mark.asyncio +async def test_breach_check_skipped_in_offline_mode(monkeypatch): + monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") + password = "correct horse battery staple" + result = await validate_password_not_breached( + password=password, + general_settings={}, + client=_client_returning(f"{_sha1_upper(password)[5:]}:42"), + ) + assert result is None + + @pytest.mark.asyncio async def test_rejects_breached_password(): password = "correct horse battery staple" diff --git a/tests/unit/proxy/auth/test_search_tool_deny_by_default.py b/tests/unit/proxy/auth/test_search_tool_deny_by_default.py index cd285e0386f..da036795889 100644 --- a/tests/unit/proxy/auth/test_search_tool_deny_by_default.py +++ b/tests/unit/proxy/auth/test_search_tool_deny_by_default.py @@ -1,8 +1,9 @@ from collections.abc import Mapping -from typing import Final, TypedDict +from typing import Final, TypedDict, cast from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy import proxy_server @@ -78,6 +79,14 @@ def _team_loader(team_tools: list[str] | None | str) -> TeamObjectLoader: return load +def _missing_dashboard_team_loader() -> tuple[TeamObjectLoader, AsyncMock]: + missing_team: Final = HTTPException( + status_code=404, detail="Team doesn't exist in db. Team=litellm-dashboard" + ) + load_team: Final = AsyncMock(side_effect=missing_team) + return cast(TeamObjectLoader, load_team), load_team + + async def _no_team() -> None: return None @@ -194,9 +203,6 @@ async def test_search_tool_access_is_unchanged_without_search_tool_deny_by_defau id="proxy admin virtual key not exempt", ), pytest.param(CallerFields(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS), None, None, None, id="master key exempt"), - pytest.param( - CallerFields(team_id=UI_TEAM_ID, virtual_key=False), None, None, None, id="dashboard session exempt" - ), ], ) @pytest.mark.asyncio @@ -248,6 +254,79 @@ async def test_search_tool_deny_by_default_does_not_load_the_user_when_off( get_user_object.assert_not_awaited() +@pytest.mark.parametrize( + "general_settings", + [{}, {"search_tool_deny_by_default": False}], + ids=["omitted", "false"], +) +@pytest.mark.parametrize( + "user_role", + [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.INTERNAL_USER], + ids=["admin", "internal-user"], +) +@pytest.mark.parametrize( + "key_tools, expected", + [ + pytest.param(None, None, id="unrestricted-key"), + pytest.param(["search-a"], None, id="key-grants-search-a"), + pytest.param(["search-b"], _LEGACY, id="key-excludes-search-a"), + ], +) +@pytest.mark.asyncio +async def test_dashboard_search_tool_grants_without_deny_by_default( + general_settings: Mapping[str, object], + user_role: LitellmUserRoles, + key_tools: list[str] | None, + expected: ProxyErrorTypes | None, +) -> None: + caller: Final = _caller( + virtual_key=False, + key_tools=key_tools, + team_id=UI_TEAM_ID, + user_role=user_role, + ) + load_team, load_team_mock = _missing_dashboard_team_loader() + assert await _denied_by(general_settings, caller, load_team) == expected + load_team_mock.assert_not_awaited() + + +@pytest.mark.parametrize( + "user_role", + [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], + ids=["admin", "admin-view-only"], +) +@pytest.mark.asyncio +async def test_admin_dashboard_search_tool_grants_with_deny_by_default(user_role: LitellmUserRoles) -> None: + caller: Final = _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=user_role) + load_team, load_team_mock = _missing_dashboard_team_loader() + assert await _denied_by(_DENY_ON, caller, load_team) is None + load_team_mock.assert_not_awaited() + + +@pytest.mark.parametrize( + "user_id, user_tools, expected", + [ + pytest.param("user-1", ["search-a"], None, id="user-grants-search-a"), + pytest.param("user-1", ["search-b"], _USER, id="user-grants-another-tool"), + pytest.param("user-1", None, _USER, id="user-grant-null"), + pytest.param("user-2", None, _USER, id="no-user-row"), + ], +) +@pytest.mark.asyncio +async def test_internal_dashboard_search_tool_requires_the_user_grant_with_deny_by_default( + cache: UserApiKeyCache, + user_id: str, + user_tools: list[str] | None, + expected: ProxyErrorTypes | None, +) -> None: + if user_id == "user-1": + _cache_user(cache, user_tools) + caller: Final = _caller(virtual_key=False, team_id=UI_TEAM_ID, user_id=user_id) + load_team, load_team_mock = _missing_dashboard_team_loader() + assert await _denied_by(_DENY_ON, caller, load_team) == expected + load_team_mock.assert_not_awaited() + + @pytest.mark.parametrize( "general_settings, caller, expected", [ @@ -257,6 +336,24 @@ async def test_search_tool_deny_by_default_does_not_load_the_user_when_off( pytest.param(_DENY_ON, _caller(virtual_key=False), _USER, id="keyless user"), pytest.param(_DENY_ON, _caller(user_role=LitellmUserRoles.PROXY_ADMIN), _KEY, id="proxy admin key"), pytest.param(_DENY_ON, _caller(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS), None, id="master key"), + pytest.param( + _DENY_ON, + _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.INTERNAL_USER), + _USER, + id="internal dashboard user", + ), + pytest.param( + _DENY_ON, + _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.PROXY_ADMIN), + None, + id="admin dashboard", + ), + pytest.param( + _DENY_ON, + _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), + None, + id="admin view-only dashboard", + ), ], ) def test_unregistered_search_fallback_follows_search_tool_deny_by_default( @@ -268,3 +365,12 @@ def test_unregistered_search_fallback_follows_search_tool_deny_by_default( with pytest.raises(ProxyException) as exc_info: check_unregistered_search_fallback(caller, general_settings) assert (exc_info.value.type, exc_info.value.code) == (expected, "403") + + +@pytest.mark.asyncio +async def test_search_preserves_missing_real_team_error_without_deny_by_default() -> None: + missing_team: Final = HTTPException(status_code=404, detail="Team doesn't exist in db. Team=team-1") + load_team: Final = AsyncMock(side_effect=missing_team) + with pytest.raises(HTTPException) as exc_info: + await can_caller_call_search_tool("search-a", _caller(team_id="team-1"), {}, load_team) + assert exc_info.value is missing_team diff --git a/tests/unit/proxy/auth/test_team_member_budget.py b/tests/unit/proxy/auth/test_team_member_budget.py index b38a953d189..6055e4f0356 100644 --- a/tests/unit/proxy/auth/test_team_member_budget.py +++ b/tests/unit/proxy/auth/test_team_member_budget.py @@ -3,8 +3,11 @@ Unit tests for team member budget checks in common_checks. These tests verify the team member budget enforcement without requiring a proxy server. """ -import pytest +from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch + +import pytest from fastapi import Request import litellm @@ -15,7 +18,10 @@ from litellm.proxy._types import ( LiteLLM_UserTable, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_checks import common_checks, get_team_membership +from litellm.proxy.auth.auth_checks import ( + common_checks, + resolve_effective_team_member_budget, +) @pytest.mark.asyncio @@ -436,3 +442,48 @@ async def test_team_member_budget_check_personal_key_not_team(): # Should pass and get_team_membership should not be called assert result is True mock_get_team_membership.assert_not_called() + + +def test_resolve_effective_team_member_budget_preserves_member_and_default_precedence() -> None: + now: Final = datetime(2030, 1, 1, tzinfo=timezone.utc) + member_budget: Final = LiteLLM_BudgetTable( + max_budget=100, + temp_budget_increase=10, + temp_budget_expiry=datetime(2030, 1, 2, tzinfo=timezone.utc), + ) + expired_member_budget: Final = LiteLLM_BudgetTable( + max_budget=100, + temp_budget_increase=10, + temp_budget_expiry=datetime(2029, 12, 31, tzinfo=timezone.utc), + ) + active_inherited_increase: Final = LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=10, + temp_budget_expiry=datetime(2030, 1, 2, tzinfo=timezone.utc), + ) + expired_inherited_increase: Final = LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=10, + temp_budget_expiry=datetime(2029, 12, 31, tzinfo=timezone.utc), + ) + + assert resolve_effective_team_member_budget(member_budget, None, now) == 110 + assert resolve_effective_team_member_budget(expired_member_budget, None, now) == 100 + assert resolve_effective_team_member_budget( + active_inherited_increase, + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 60 + assert resolve_effective_team_member_budget( + expired_inherited_increase, + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 50 + assert resolve_effective_team_member_budget( + LiteLLM_BudgetTable(max_budget=0), + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 0 + assert resolve_effective_team_member_budget(None, LiteLLM_BudgetTable(max_budget=50), now) == 50 + assert resolve_effective_team_member_budget(None, LiteLLM_BudgetTable(max_budget=0), now) is None + assert resolve_effective_team_member_budget(None, None, now) is None diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 4cdd7fefd28..2d9c781b0d1 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -7724,7 +7724,8 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() @pytest.mark.asyncio -async def test_user_api_key_auth_authenticates_before_raising_malformed_body_error(): +@pytest.mark.parametrize("body", [b'{}{"model": "gpt-4o"}', b"[]", b"null", b"[1e400]"]) +async def test_user_api_key_auth_authenticates_before_raising_malformed_body_error(body: bytes) -> None: """Regression (LIT-4780): a body that fails to parse must still be authenticated first, so the rejected request's trace carries the caller's key / team / user identity instead of an anonymous root span. The parse error is re-raised @@ -7742,10 +7743,11 @@ async def test_user_api_key_auth_authenticates_before_raising_malformed_body_err "headers": [(b"content-type", b"application/json")], "method": "POST", "path": "/chat/completions", + "query_string": b"probe=1", } ) request._url = URL(url="/chat/completions") - request._body = b'{}{"model": "gpt-4o"}' + request._body = body attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} @@ -11688,6 +11690,124 @@ def _alias_request(route: str, data: dict, content_type: str = "application/json return request +def _viewer_inference_key(monkeypatch: pytest.MonkeyPatch, team_id: str | None, *, spend: float = 0.25) -> str: + from litellm.proxy import proxy_server + from litellm.proxy.utils import ProxyLogging, hash_token + + key: Final = "sk-viewer-inference" + hashed_key: Final = hash_token(key) + cache: Final = UserApiKeyCache() + cache.set_cache( + key=hashed_key, + value=UserAPIKeyAuth( + token=hashed_key, + user_id="viewer", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + team_id=team_id, + models=["gpt-6.1-sol"], + max_budget=1.0, + spend=spend, + ), + ) + cache.set_cache( + key="viewer", + value=LiteLLM_UserTable( + user_id="viewer", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + user_email="viewer@example.com", + models=["gpt-6.1-sol"], + ), + ) + if team_id is not None: + cache.set_cache( + key=team_membership_reservation_cache_key(user_id="viewer", team_id=team_id), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) + for name, value in { + **_proxy_attrs_for_centralized_checks(), + "prisma_client": MagicMock(), + "user_api_key_cache": cache, + "spend_counter_cache": DualCache(), + "proxy_logging_obj": ProxyLogging(user_api_key_cache=cache), + "llm_model_list": [], + }.items(): + monkeypatch.setattr(proxy_server, name, value) + return key + + +@pytest.mark.asyncio +@pytest.mark.parametrize("team_id", ["litellm-dashboard", None], ids=["dashboard-session", "virtual-key"]) +@pytest.mark.parametrize("route", ["/chat/completions", "/v1/chat/completions"]) +async def test_viewer_inference_uses_the_normal_key_authorization_path( + monkeypatch: pytest.MonkeyPatch, team_id: str | None, route: str +) -> None: + key: Final = _viewer_inference_key(monkeypatch, team_id) + request: Final = _alias_request(route, {"model": "gpt-6.1-sol", "messages": [{"role": "user", "content": "usage"}]}) + + result: Final = await user_api_key_auth(request=request, api_key=f"Bearer {key}") + + assert (result.user_id, result.user_role, result.team_id, result.models, result.max_budget) == ( + "viewer", + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + team_id, + ["gpt-6.1-sol"], + 1.0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("team_id", ["litellm-dashboard", None], ids=["dashboard-session", "virtual-key"]) +async def test_viewer_inference_rejects_a_model_outside_the_key_allowlist( + monkeypatch: pytest.MonkeyPatch, team_id: str | None +) -> None: + key: Final = _viewer_inference_key(monkeypatch, team_id) + request: Final = _alias_request("/chat/completions", {"model": "blocked-model"}) + + with pytest.raises(ProxyException) as failure: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + + assert (failure.value.code, failure.value.type, failure.value.param) == ( + "403", + ProxyErrorTypes.key_model_access_denied, + "model", + ) + assert failure.value.message == ( + "The requested model 'blocked-model' is not available for this API key, or the model name is invalid. " + "Check the models available to you and try again." + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("team_id", ["litellm-dashboard", None], ids=["dashboard-session", "virtual-key"]) +@pytest.mark.parametrize("spend", [1.0, 1.25], ids=["at-limit", "over-limit"]) +async def test_viewer_inference_rejects_an_exhausted_key_budget( + monkeypatch: pytest.MonkeyPatch, team_id: str | None, spend: float +) -> None: + key: Final = _viewer_inference_key(monkeypatch, team_id, spend=spend) + request: Final = _alias_request("/chat/completions", {"model": "gpt-6.1-sol"}) + + with pytest.raises(ProxyException) as failure: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + + assert failure.value.type == ProxyErrorTypes.budget_exceeded + assert f"Current cost: {spend}, Max budget: 1.0" in failure.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("team_id", ["litellm-dashboard", None], ids=["dashboard-session", "virtual-key"]) +async def test_viewer_inference_access_does_not_allow_management_writes( + monkeypatch: pytest.MonkeyPatch, team_id: str | None +) -> None: + key: Final = _viewer_inference_key(monkeypatch, team_id) + request: Final = _alias_request("/team/new", {"team_alias": "forbidden"}) + + with pytest.raises(ProxyException) as failure: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + + assert failure.value.code == "403" + assert "proxy_admin_viewer" in failure.value.message + + async def _enforce_alias_access(token: UserAPIKeyAuth, data: dict, route: str, request, router: litellm.Router): from litellm.proxy.auth.user_api_key_auth import enforce_key_and_fallback_model_access @@ -12480,11 +12600,14 @@ async def test_lens_oauth2_validates_identity_without_inference_user_field( patch("litellm.proxy.proxy_server.master_key", "sk-fixture-master"), patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch.dict(os.environ, { - "OAUTH_TOKEN_INFO_ENDPOINT": "https://oauth.fixture.test/introspect", - "OAUTH_CLIENT_ID": "fixture-client", - "OAUTH_CLIENT_SECRET": "fixture-client-secret", - }), + patch.dict( + os.environ, + { + "OAUTH_TOKEN_INFO_ENDPOINT": "https://oauth.fixture.test/introspect", + "OAUTH_CLIENT_ID": "fixture-client", + "OAUTH_CLIENT_SECRET": "fixture-client-secret", + }, + ), patch("httpx.AsyncClient.send", side_effect=token_info) as network, ): if active: @@ -12611,6 +12734,4 @@ def test_should_skip_budget_checks_spares_only_zero_spend_mcp_discovery( ) -> None: from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks - assert ( - _should_skip_budget_checks(request_data=request_data, route=route, request=None, llm_router=None) is expected - ) + assert _should_skip_budget_checks(request_data=request_data, route=route, request=None, llm_router=None) is expected diff --git a/tests/unit/proxy/client/cli/test_encryption_commands.py b/tests/unit/proxy/client/cli/test_encryption_commands.py index 3a86eb82593..6685d29fb7e 100644 --- a/tests/unit/proxy/client/cli/test_encryption_commands.py +++ b/tests/unit/proxy/client/cli/test_encryption_commands.py @@ -5,6 +5,8 @@ check vs POST migrate, dry-run param) and its residual-state messaging without a live proxy. """ +from typing import Final + import pytest from click.testing import CliRunner @@ -69,3 +71,20 @@ def test_migrate_reports_residual_legacy(runner): result = runner.invoke(cli_main.cli, ["encryption", "migrate", "--check"]) assert result.exit_code == 0, result.output assert "Residual legacy values remaining: 3" in result.output + + +@pytest.mark.parametrize("residual", [-1, float("nan")]) +def test_migrate_prints_no_residual_status_for_a_residual_below_zero_or_nan(runner, residual): + _FakeHTTPClient.response = {"status": "success", "report": {"residual_legacy": residual}} + result: Final = runner.invoke(cli_main.cli, ["encryption", "migrate", "--check"]) + assert result.exit_code == 0, result.output + assert "No legacy values remaining" not in result.output + assert "Residual legacy values remaining" not in result.output + + +@pytest.mark.parametrize("report", [None, "done", {"residual_legacy": "3"}]) +def test_migrate_malformed_report_fails(runner, report): + _FakeHTTPClient.response = {"status": "success", "report": report} + result: Final = runner.invoke(cli_main.cli, ["encryption", "migrate", "--check"]) + assert result.exit_code != 0 + assert "Unexpected" in result.output diff --git a/tests/unit/proxy/client/cli/test_pi.py b/tests/unit/proxy/client/cli/test_pi.py index ff4dc25cd0c..e6669564358 100644 --- a/tests/unit/proxy/client/cli/test_pi.py +++ b/tests/unit/proxy/client/cli/test_pi.py @@ -195,7 +195,12 @@ class TestProviderBlock: assert block == { "baseUrl": "http://localhost:4000/v1", "api": "openai-completions", - "compat": {"supportsStore": False, "supportsLongCacheRetention": False}, + "compat": { + "supportsStore": False, + "supportsLongCacheRetention": False, + "sendSessionAffinityHeaders": True, + "sessionAffinityFormat": "openrouter", + }, "apiKey": "$LITELLM_PROXY_API_KEY", "models": [{"id": "m-1"}, {"id": "m-2"}], } diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 33b8085e082..d7fc844171d 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -4259,6 +4259,17 @@ class TestCreatorTurnedOffMessageLogging: assert await instance._creator_turned_off_message_logging(self._job()) is expected + @pytest.mark.asyncio + async def test_a_job_without_creator_fields_falls_back_to_the_proxy_setting(self, monkeypatch): + from types import SimpleNamespace + + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + instance: Final = self._instance() + + assert await instance._creator_turned_off_message_logging(SimpleNamespace(id="job-1")) is False + @pytest.mark.asyncio async def test_a_team_that_turned_off_message_logging_covers_its_keys(self, monkeypatch): from types import SimpleNamespace diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index cd5820fa2fd..354feb0bab1 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -64,6 +64,100 @@ def _starlette_request( return Request(scope, receive) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [b"[]", b'[["model", "test"]]', b"null", b'"s"', b"1", b"true", b"false"], +) +@pytest.mark.parametrize("path", ["/v1/messages", "/vertex_ai/v1/messages", "/lens/datasets"]) +async def test_read_request_body_rejects_non_object_json(body: bytes, path: str) -> None: + request: Final = _starlette_request(body, "application/json", path) + + for _ in range(2): + with pytest.raises(ProxyException) as exc_info: + await read_request_body(request) + assert exc_info.value.code == "400" + assert exc_info.value.type == "invalid_request_error" + assert exc_info.value.param == "request_body" + assert exc_info.value.message == "Invalid JSON payload: expected an object" + assert _safe_get_request_parsed_body(request) is None + assert await request.body() == body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [b"[1e400]", b"1e400", b"[NaN]"]) +async def test_read_request_body_rejects_non_objects_from_json_fallback(body: bytes) -> None: + with pytest.raises(ProxyException) as exc_info: + await read_request_body(_starlette_request(body, "application/json")) + assert exc_info.value.code == "400" + assert exc_info.value.message == "Invalid JSON payload: expected an object" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body, invalid_unicode", + [ + (b"[]", False), + (b"null", False), + (b'"value"', False), + (b"[1e400]", False), + (b'"\\ud800"', True), + (b'["\\ud800"]', True), + (b'{"value":"\\ud800"}', True), + ], +) +@pytest.mark.parametrize("dispatch", ["custom", "provider", "lens"]) +async def test_pass_through_body_validation_preserves_bytes(dispatch: str, body: bytes, invalid_unicode: bool) -> None: + from fastapi.routing import APIRoute + from litellm.proxy.lens.adapter import lens_request + from litellm.proxy.lens.buffering import LensRoute + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import vertex_proxy_route + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import create_pass_through_route + + path: Final = {"custom": "/v1/messages", "provider": "/vertex_ai/predict", "lens": "/lens/datasets"}[dispatch] + route: Final = { + "custom": APIRoute(path, create_pass_through_route(path, "https://upstream.test"), methods=["POST"]), + "provider": APIRoute("/vertex_ai/{endpoint:path}", vertex_proxy_route, methods=["POST"]), + "lens": LensRoute("/lens/{path:path}", lens_request, methods=["POST"]), + }[dispatch] + incoming: Final = _starlette_request(body, "application/json", path) + _, dispatched_scope = route.matches(incoming.scope) + request: Final = Request({**incoming.scope, **dispatched_scope}, incoming.receive) + + for _ in range(2): + if invalid_unicode: + with pytest.raises(ProxyException) as exc_info: + await read_request_body(request) + assert exc_info.value.code == "400" + else: + assert await read_request_body(request) == {} + assert await request.body() == body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("optional", [False, True]) +@pytest.mark.parametrize("body", [b"null", b" \nnull\t", b"[]", b'"value"']) +async def test_null_body_respects_the_dispatched_body_requirement(optional: bool, body: bytes) -> None: + from fastapi.routing import APIRoute + from litellm.proxy.management_endpoints.key_management_endpoints import generate_key_fn + from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_end_users + + endpoint: Final = global_spend_end_users if optional else generate_key_fn + route: Final = APIRoute("/audit", endpoint, methods=["POST"]) + incoming: Final = _starlette_request(body, "application/json", "/audit") + _, dispatched_scope = route.matches(incoming.scope) + request: Final = Request({**incoming.scope, **dispatched_scope}, incoming.receive) + + if optional and body.strip() == b"null": + assert await read_request_body(request) == {} + assert await request.json() is None + else: + with pytest.raises(ProxyException) as exc_info: + await read_request_body(request) + assert exc_info.value.code == "400" + assert await request.body() == body + + @pytest.mark.asyncio async def test_read_request_body_marks_body_received_once_with_its_size(monkeypatch: pytest.MonkeyPatch): events: list[tuple[str, dict[str, str | int]]] = [] # mutable-ok: recorder for the injected phase_event double @@ -131,14 +225,20 @@ async def test_read_raw_json_body_is_none_for_form_bodies(): @pytest.mark.asyncio @pytest.mark.parametrize("content_type", ["application/x-protobuf", "application/protobuf; charset=binary"]) -async def test_protobuf_body_is_not_parsed_as_json(content_type): - # OTLP trace exports (POST /v1/traces) are binary protobuf; arbitrary bytes like these - # used to hit the JSON surrogate-repair path and fail auth with a 400. - body = b"\n\xa2\x01\n\x1c\n\x0cservice.name\x12\x0c\n\nswarm\xed\xa0\x80\xff" - request = _starlette_request(body, content_type) +@pytest.mark.parametrize( + "body", + [ + b"\n\xa2\x01\n\x1c\n\x0cservice.name\x12\x0c\n\nswarm\xed\xa0\x80\xff", + pytest.param(b"100000000", id="protobuf-also-json-number"), + ], +) +async def test_protobuf_body_is_not_parsed_as_json(content_type: str, body: bytes) -> None: + request: Final = _starlette_request(body, content_type) assert await read_request_body(request) == {} - assert await request.body() == body # body is still readable by the endpoint + assert await read_request_body(request) == {} + assert await read_raw_json_body(request) == body + assert await request.body() == body @pytest.mark.asyncio diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 70f7774e691..183dd7919c1 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -113,6 +113,9 @@ class MockBatcher: self.litellm_projecttable = _Table("project", self) self.litellm_endusertable = _Table("enduser", self) + def execute_raw(self, query: str, *args: object) -> None: + self.calls.append({"table": "raw", "op": "execute_raw", "query": query, "args": args}) + async def commit(self): self.committed = True return self.calls @@ -2112,6 +2115,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo assert {(call["table"], call["op"]) for call in enduser_page.calls} == {("enduser", "update_many")} assert batcher.committed is True assert {(call["table"], call["op"]) for call in batcher.calls} == { + ("raw", "execute_raw"), ("team_membership", "update_many"), ("key", "update_many"), ("org", "update_many"), @@ -3393,6 +3397,35 @@ def test_budget_cascade_carries_overage_per_tier_when_rollover_enabled( counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team_member:member-1:team-1") +def test_budget_cascade_floors_membership_total_spend_before_zeroing_spend( + rollover_enabled: None, + reset_budget_job: ResetBudgetJob, + mock_prisma_client: MockPrismaClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Membership rows written before total_spend existed hold their history only in + spend. The cascade lifts total_spend to spend for the expiring tier's rows, in the + same batch and ahead of the zero/decrement, so the reset cannot erase it.""" + _make_counter_invalidation_job(monkeypatch) + budget: Final = _budget_row(budget_id="budget-legacy", budget_duration="7d", max_budget=10.0) + mock_prisma_client.data["budget"] = [budget] + membership: Final = type( + "Membership", + (), + {"user_id": "member-1", "team_id": "team-1", "spend": 15.0, "budget_id": "budget-legacy"}, + ) + mock_prisma_client.db.litellm_teammembership.set_find_many_results([membership]) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + writes: Final = mock_prisma_client.db.batch_calls + repair_index: Final = next(i for i, call in enumerate(writes) if call["op"] == "execute_raw") + first_membership_reset: Final = next(i for i, call in enumerate(writes) if call["table"] == "team_membership") + assert repair_index < first_membership_reset, writes + assert writes[repair_index]["args"] == (["budget-legacy"],) + assert mock_prisma_client.db.batchers[-1].committed + + def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch ): diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index f30606ac2f1..24c12e06c9c 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyTagSpendTransaction @@ -769,3 +768,18 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_ assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False assert redis.rows == [] + + +@pytest.mark.asyncio +async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails(): + from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue + + redis_cache = MagicMock() + redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + await RedisUpdateBuffer(redis_cache=redis_cache).store_tool_discoveries_in_redis(queue) + + assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2}] diff --git a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index e400ad16e84..0a015255f2e 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -5,9 +5,9 @@ Unit tests for ToolDiscoveryQueue. import pytest - from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, + aggregate_tool_discoveries, ) @@ -16,57 +16,50 @@ def queue(): return ToolDiscoveryQueue() -def test_add_single_tool(queue): +def test_a_tool_seen_once_flushes_with_call_count_one(queue): queue.add_update({"tool_name": "my_tool", "origin": "user_defined"}) - items = queue.flush() - assert len(items) == 1 - assert items[0]["tool_name"] == "my_tool" - assert items[0]["origin"] == "user_defined" + assert queue.flush() == [{"tool_name": "my_tool", "origin": "user_defined", "call_count": 1}] -def test_deduplication_same_name(queue): - """Adding the same tool_name twice should only keep the first.""" - queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"}) - queue.add_update({"tool_name": "tool_a", "origin": "user_defined"}) - items = queue.flush() - assert len(items) == 1 - assert items[0]["origin"] == "mcp_server" # first wins +def test_repeated_sightings_fold_into_one_row_counting_every_call(queue): + queue.add_update({"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice"}) + queue.add_update({"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob"}) + queue.add_update({"tool_name": "tool_a"}) + assert queue.flush() == [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice", "call_count": 3}] -def test_deduplication_different_names(queue): +def test_distinct_tools_each_get_their_own_row(queue): queue.add_update({"tool_name": "tool_a"}) queue.add_update({"tool_name": "tool_b"}) - items = queue.flush() - assert len(items) == 2 - names = {i["tool_name"] for i in items} - assert names == {"tool_a", "tool_b"} - - -def test_flush_clears_pending(queue): - queue.add_update({"tool_name": "tool_x"}) - items1 = queue.flush() - assert len(items1) == 1 - items2 = queue.flush() - assert len(items2) == 0 - - -def test_seen_names_reset_after_flush(queue): - """Seen-set is cleared on flush so the same tool can re-enter the next cycle.""" queue.add_update({"tool_name": "tool_a"}) - queue.flush() - queue.add_update({"tool_name": "tool_a"}) # same tool, new cycle - items = queue.flush() - assert len(items) == 1 - assert items[0]["tool_name"] == "tool_a" + assert queue.flush() == [ + {"tool_name": "tool_a", "call_count": 2}, + {"tool_name": "tool_b", "call_count": 1}, + ] + + +def test_flush_clears_pending_and_the_next_cycle_counts_from_zero(queue): + queue.add_update({"tool_name": "tool_x"}) + queue.add_update({"tool_name": "tool_x"}) + assert queue.flush() == [{"tool_name": "tool_x", "call_count": 2}] + assert queue.flush() == [] + queue.add_update({"tool_name": "tool_x"}) + assert queue.flush() == [{"tool_name": "tool_x", "call_count": 1}] def test_empty_tool_name_ignored(queue): queue.add_update({"tool_name": ""}) queue.add_update({"tool_name": None}) # type: ignore[arg-type] - items = queue.flush() - assert len(items) == 0 + assert queue.flush() == [] -def test_flush_returns_list(queue): - result = queue.flush() - assert isinstance(result, list) +def test_aggregate_sums_call_counts_across_workers_flushes(): + worker_a = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 3}] + worker_b = [ + {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 2}, + {"tool_name": "search", "call_count": 1}, + ] + assert aggregate_tool_discoveries([*worker_a, *worker_b]) == ( + {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5}, + {"tool_name": "search", "call_count": 1}, + ) diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 7437042e4bd..a42b6404278 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -19,7 +19,7 @@ from redis.exceptions import DataError import litellm from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes -from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_UPDATE_BUFFER_KEY, SPEND_ROLLUP_LOCK_TIMEOUT_MS +from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import ( _TEAM_ADVISORY_LOCK_SQL, @@ -29,13 +29,13 @@ from litellm.proxy.db.db_spend_update_writer import ( _SpendTableName, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue -from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, ) from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL +from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -1128,7 +1128,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))" in statement assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement - assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement + assert ( + 'total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend)' + " + EXCLUDED.total_spend" + ) in statement @pytest.mark.asyncio @@ -2158,46 +2161,122 @@ async def test_commit_spend_updates_to_db_increments_key_total_spend_alongside_s @pytest.mark.asyncio -async def test_update_database_creates_single_task(): +async def test_shutdown_drain_finishes_the_rollups_update_database_started(): """ - Test that update_database() fires exactly 1 asyncio.create_task() call - (the batched task) instead of the previous 11. + Regression for LIT-6027. The key, team and daily rollups ran in an unreferenced asyncio task, + so the shutdown sequence (flush, drain_logging_tasks, flush) never waited for them and spend + logs landed without matching daily rows. Once update_database returns, drain_logging_tasks + must leave every rollup for that request in its queue. """ - db_writer = DBSpendUpdateWriter() - - # Mock all helpers so nothing real runs + db_writer: Final = DBSpendUpdateWriter() db_writer._insert_spend_log_to_db = AsyncMock() - db_writer._batch_database_updates = AsyncMock() - + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + fake_payload: Final = { + "request_id": "req-lit-6027", + "api_key": "hashed-key", + "user": "test-user", + "team_id": "test-team", + "agent_id": None, + "startTime": "2024-01-01T00:00:00", + "endTime": "2024-01-01T00:00:01", + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "model_group": "gpt-4o-mini", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.0, + "metadata": '{"usage_object": {}}', + } with ( patch("litellm.proxy.proxy_server.disable_spend_logs", False), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), - patch("litellm.proxy.db.db_spend_update_writer.asyncio.create_task") as mock_create_task, + patch("litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", return_value=fake_payload), ): await db_writer.update_database( token="test-token", user_id="test-user", - end_user_id="test-end-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4o-mini", "custom_llm_provider": "openai"}, + completion_response=MagicMock(), start_time=datetime.now(), end_time=datetime.now(), - team_id="test-team", - org_id="test-org", - completion_response=MagicMock(), response_cost=0.1, - kwargs={"model": "gpt-4", "custom_llm_provider": "openai"}, ) + await drain_logging_tasks() - # Exactly 1 create_task call (the batch), not 11 - assert mock_create_task.call_count == 1 + spend: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert spend["key_list_transactions"] == {"test-token": 0.1} + assert spend["team_list_transactions"] == {"test-team": 0.1} + daily_user: Final = ( + await db_writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_team: Final = ( + await db_writer.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + assert daily_user is not None and daily_team is not None + (user_row,) = daily_user.values() + (team_row,) = daily_team.values() + assert (user_row["api_requests"], user_row["spend"]) == (1, 0.1) + assert (team_row["api_requests"], team_row["spend"]) == (1, 0.1) @pytest.mark.asyncio -async def test_batch_database_updates_isolation_on_failure(): +async def test_shutdown_drain_waits_for_an_in_flight_rollup_enqueue(): + """ + Regression for LIT-6027. A rollup enqueue that is still in flight when the shutdown drain + starts (here: blocked on a full spend queue) must be finished by the time drain_logging_tasks + returns, so the second shutdown flush commits it. + """ + enqueue_started: Final = asyncio.Event() + release_enqueue: Final = asyncio.Event() + + class _BlockedSpendQueue(SpendUpdateQueue): + async def add_update(self, update: SpendUpdateQueueItem): + enqueue_started.set() + await release_enqueue.wait() + await super().add_update(update) + + db_writer: Final = DBSpendUpdateWriter() + db_writer.spend_update_queue = _BlockedSpendQueue() + db_writer._insert_spend_log_to_db = AsyncMock() + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", True), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + ): + await db_writer.update_database( + token="t1", + user_id=None, + end_user_id=None, + team_id="team-1", + org_id=None, + kwargs={"model": "gpt-4o-mini", "custom_llm_provider": "openai"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + await enqueue_started.wait() + drain: Final = asyncio.create_task(drain_logging_tasks()) + release_enqueue.set() + await drain + spend: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + + assert spend["key_list_transactions"] == {"t1": 0.1} + assert spend["team_list_transactions"] == {"team-1": 0.1} + + +@pytest.mark.asyncio +async def test_batch_database_updates_isolation_on_failure(caplog: pytest.LogCaptureFixture): """ Test that if one helper inside _batch_database_updates raises, - all other helpers still execute. + all other helpers still execute and the failure is reported at ERROR with the request id. """ db_writer = DBSpendUpdateWriter() @@ -2217,6 +2296,7 @@ async def test_batch_database_updates_isolation_on_failure(): db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() + caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name) await db_writer._batch_database_updates( response_cost=0.1, user_id="u1", @@ -2226,7 +2306,7 @@ async def test_batch_database_updates_isolation_on_failure(): end_user_id="eu1", prisma_client=MagicMock(), litellm_proxy_budget_name="budget", - payload={"key": "value"}, + payload={"request_id": "req-iso", "key": "value"}, ) # _update_key_db raised, but all others should still have been called @@ -2242,6 +2322,48 @@ async def test_batch_database_updates_isolation_on_failure(): db_writer.add_spend_log_transaction_to_daily_team_transaction.assert_awaited_once() db_writer.add_spend_log_transaction_to_daily_org_transaction.assert_awaited_once() db_writer.add_spend_log_transaction_to_daily_tag_transaction.assert_awaited_once() + assert any( + "_update_key_db" in record.getMessage() and "req-iso" in record.getMessage() + for record in caplog.records + if record.levelno >= logging.ERROR + ) + + +@pytest.mark.asyncio +async def test_malformed_metadata_keeps_the_daily_rollup_row_and_is_reported(caplog: pytest.LogCaptureFixture): + """ + Regression for LIT-6027. A spend log whose metadata is not JSON used to vanish from every + daily rollup while the spend log itself persisted, with nothing above debug in the logs. + """ + writer: Final = DBSpendUpdateWriter() + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + payload: Final = { + "request_id": "req-bad-metadata", + "user": "test-user", + "startTime": "2024-01-01T00:00:00", + "endTime": "2024-01-01T00:00:01", + "api_key": "test-key", + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "model_group": "gpt-4o-mini", + "prompt_tokens": 100, + "completion_tokens": 50, + "spend": 0.05, + "metadata": "{not-json", + } + caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name) + + await writer.add_spend_log_transaction_to_daily_user_transaction( # pyright: ignore[reportUnknownMemberType] # prisma_client param is untyped on the writer + payload=payload, prisma_client=prisma + ) + + daily: Final = await writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + (row,) = daily.values() + assert (row["api_requests"], row["spend"], row["prompt_tokens"]) == (1, 0.05, 100) + assert any( + "req-bad-metadata" in record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR + ) @pytest.mark.asyncio @@ -2394,31 +2516,6 @@ class _InMemoryRedis: return [self.lists.pop(op["key"], None) for op in lpop_list] -@pytest.mark.asyncio -@pytest.mark.parametrize("lock_held_by_another_pod", [False, True]) -async def test_commit_with_redis_writes_discovered_tools_on_leader_and_non_leader(lock_held_by_another_pod: bool): - redis = _InMemoryRedis() - if lock_held_by_another_pod: - assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=DB_SPEND_UPDATE_JOB_NAME) - db_writer = DBSpendUpdateWriter(redis_cache=redis) - db_writer.pod_lock_manager.redis_cache = redis - db_writer.tool_discovery_queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) - prisma_client = MagicMock() - prisma_client.db.litellm_tooltable.upsert = AsyncMock() - - await db_writer._commit_spend_updates_to_db_with_redis( - prisma_client=prisma_client, - n_retry_times=0, - proxy_logging_obj=MagicMock(), - ) - - upserted = [c.kwargs["where"] for c in prisma_client.db.litellm_tooltable.upsert.await_args_list] - assert upserted == [{"tool_name": "get_weather"}] - assert (REDIS_UPDATE_BUFFER_KEY in redis.lists) is lock_held_by_another_pod, ( - f"spend buffer left in redis={list(redis.lists)}, only the lock holder may drain it" - ) - - @pytest.mark.asyncio async def test_commit_with_redis_requeues_all_on_db_failure(): """ @@ -2775,12 +2872,12 @@ async def test_update_database_does_not_deepcopy_on_request_path(): """ Regression for LIT-4088: copy.deepcopy must not run while the caller awaits update_database(). The deepcopy used to isolate the daily-spend helpers is - relocated into the _batch_database_updates background task, and the spend-log + relocated into the _batch_database_updates logging task, and the spend-log insert receives the payload directly (all consumers are read-only). Asserts: - zero copy.deepcopy calls happen on the awaited request path - - the batch background task still hands the daily helpers an isolated copy + - the tracked rollup task still hands the daily helpers an isolated copy (mutating the original after the task ran does not bleed into it) - the spend-log insert receives the payload on the request path with the correct content @@ -2856,16 +2953,13 @@ async def test_update_database_does_not_deepcopy_on_request_path(): response_cost=0.1, ) - # Request path is clean: nothing was deepcopied while the caller awaited. assert len(deepcopy_calls) == 0 # The spend-log insert ran inline on the request path with the real payload. assert captured_spend_log["ref"] is fake_payload assert captured_spend_log["model_at_call"] == "gpt-4" assert fake_payload["spend"] == 0.1 - - # Now let the batch background task run; the deepcopy happens here. - await asyncio.sleep(0) + await drain_logging_tasks() assert len(deepcopy_calls) >= 1 assert len(captured_batch_payloads) == 1 @@ -3267,7 +3361,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFakeDB() db.query_raw = AsyncMock(side_effect=Exception("connection reset")) - db_writer._flush_tool_discovery_queue = AsyncMock() await db_writer._commit_spend_updates_to_db_without_redis_buffer( prisma_client=_WindowSpendFakePrisma(db), @@ -3275,7 +3368,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ proxy_logging_obj=MagicMock(), ) - db_writer._flush_tool_discovery_queue.assert_called_once() requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() assert requeued == (transaction,) @@ -3351,7 +3443,6 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov db_writer = DBSpendUpdateWriter() await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")}) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure) - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3377,7 +3468,6 @@ async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent(): {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)} ) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply")) - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3405,7 +3495,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other team_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"team_id": "team-1"} await db_writer.daily_team_spend_update_queue.add_update({"team-key": team_txn}) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend") - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3416,7 +3505,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other assert _daily_upserts(db, "LiteLLM_DailyUserSpend") == [] (team_upsert,) = _daily_upserts(db, "LiteLLM_DailyTeamSpend") assert _row_values(team_upsert, "team_id") == ["team-1"] - db_writer._flush_tool_discovery_queue.assert_called_once() db.failing_table = None await db_writer._commit_spend_updates_to_db_without_redis_buffer( @@ -3643,7 +3731,6 @@ async def test_failed_window_spend_commit_is_requeued_only_when_the_rows_are_pro ) await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFailureDB(failure) - db_writer._flush_tool_discovery_queue = AsyncMock() await db_writer._commit_spend_updates_to_db_without_redis_buffer( prisma_client=_WindowSpendFakePrisma(db), diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index c9df665741d..76898b0eeab 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -3,22 +3,28 @@ Unit tests for tool_registry_writer.py — uses a mock prisma client that exposes litellm_tooltable.upsert / find_many / find_unique. """ +import json from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock import pytest -from prisma.errors import PrismaError - +from prisma.errors import PrismaError, RawQueryError +from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_TOOL_DISCOVERY_BUFFER_KEY, TOOL_DISCOVERY_COMMIT_JOB_NAME +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue from litellm.proxy.db.tool_registry_writer import ( ToolPolicyRegistry, batch_upsert_tools, + commit_tool_discoveries, get_tool, get_tool_policy_registry, get_tools_by_names, list_tools, update_tool_policy, ) +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation def _mock_row(**kwargs): @@ -62,6 +68,7 @@ def _make_prisma( prisma = MagicMock() prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) + prisma.db.execute_raw = AsyncMock(return_value=1) prisma.db.litellm_tooltable.find_many = AsyncMock( return_value=find_many_rows if find_many_rows is not None else [] ) @@ -73,28 +80,32 @@ def _make_prisma( return prisma +def _statement(prisma) -> tuple[str, tuple[object, ...]]: + prisma.db.execute_raw.assert_awaited_once() + sql, *params = prisma.db.execute_raw.call_args.args + return sql, tuple(params) + + @pytest.mark.asyncio -async def test_batch_upsert_tools_calls_upsert(): +async def test_batch_upsert_tools_inserts_a_new_tool_with_the_batch_call_count_and_adds_it_on_conflict(): prisma = _make_prisma() - items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] - await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.upsert.assert_awaited_once() - call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs - assert call_kw["where"] == {"tool_name": "tool_a"} - assert call_kw["data"]["create"]["tool_name"] == "tool_a" - assert call_kw["data"]["create"]["origin"] == "mcp_server" - assert call_kw["data"]["create"]["input_policy"] == "untrusted" - assert call_kw["data"]["create"]["output_policy"] == "untrusted" - assert call_kw["data"]["create"]["call_count"] == 1 - assert call_kw["data"]["update"]["call_count"] == {"increment": 1} - assert "updated_at" in call_kw["data"]["update"] + await batch_upsert_tools( + prisma, [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}] + ) + sql, params = _statement(prisma) + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') + assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql + assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql + assert '"input_policy"' not in sql, "policies come from the table defaults and are never overwritten" + assert params[1:] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None) + prisma.db.litellm_tooltable.upsert.assert_not_awaited() @pytest.mark.asyncio async def test_batch_upsert_tools_empty_list(): prisma = _make_prisma() await batch_upsert_tools(prisma, []) - prisma.db.litellm_tooltable.upsert.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio @@ -102,21 +113,22 @@ async def test_batch_upsert_tools_skips_empty_names(): prisma = _make_prisma() items = [{"tool_name": "", "origin": None}, {"tool_name": None}] # type: ignore[list-item] await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.upsert.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio -async def test_batch_upsert_multiple_tools_calls_upsert_per_tool(): +async def test_batch_upsert_many_tools_is_one_statement_with_one_row_per_distinct_tool(): prisma = _make_prisma() items = [ - {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}, - {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"}, + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 2}, + {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice", "call_count": 1}, + {"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob", "call_count": 5}, ] await batch_upsert_tools(prisma, items) - assert prisma.db.litellm_tooltable.upsert.await_count == 2 - calls = prisma.db.litellm_tooltable.upsert.call_args_list - assert calls[0].kwargs["where"]["tool_name"] == "tool_a" - assert calls[1].kwargs["where"]["tool_name"] == "tool_b" + sql, params = _statement(prisma) + assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql + assert params[1:10] == ("tool_a", "mcp_server", 7, "system", "system", None, None, None, None) + assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", None, None, None, None) @pytest.mark.asyncio @@ -264,20 +276,14 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): assert call_kw["where"] == {"tool_name": "my_tool"} assert call_kw["data"]["update"]["input_policy"] == "blocked" assert call_kw["data"]["update"]["updated_by"] == "admin" - prisma.db.litellm_tooltable.find_unique.assert_awaited_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio async def test_get_tools_by_names_returns_policy_map(): rows = [ - _mock_row( - tool_name="tool_a", input_policy="trusted", output_policy="untrusted" - ), - _mock_row( - tool_name="tool_b", input_policy="blocked", output_policy="untrusted" - ), + _mock_row(tool_name="tool_a", input_policy="trusted", output_policy="untrusted"), + _mock_row(tool_name="tool_b", input_policy="blocked", output_policy="untrusted"), ] prisma = _make_prisma(find_many_rows=rows) result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) @@ -285,9 +291,7 @@ async def test_get_tools_by_names_returns_policy_map(): "tool_a": ("trusted", "untrusted"), "tool_b": ("blocked", "untrusted"), } - prisma.db.litellm_tooltable.find_many.assert_awaited_once_with( - where={"tool_name": {"in": ["tool_a", "tool_b"]}} - ) + prisma.db.litellm_tooltable.find_many.assert_awaited_once_with(where={"tool_name": {"in": ["tool_a", "tool_b"]}}) @pytest.mark.asyncio @@ -382,12 +386,8 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): return [] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock( - side_effect=_flaky_find_many - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(side_effect=_flaky_find_many) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(return_value=[]) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -397,10 +397,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): assert len(invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "sync_tool_policy_from_db_tools_lookup_failure" - ) + assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_tools_lookup_failure" assert registry.is_initialized() @@ -421,9 +418,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( - side_effect=_flaky_perms_find_many - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(side_effect=_flaky_perms_find_many) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -433,7 +428,142 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() assert len(perms_invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "sync_tool_policy_from_db_perms_lookup_failure" + assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_perms_lookup_failure" + + +class _InMemoryRedis: + def __init__(self) -> None: + self.store: dict[str, str] = {} + self.lists: dict[str, list[object]] = {} + + async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None: + if nx and key in self.store: + return None + self.store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs: object) -> str | None: + raw = self.store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.store.pop(key, None) is not None else 0 + + async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]: + for op in rpush_list: + self.lists.setdefault(op["key"], []).extend(op["values"]) + return [len(self.lists[op["key"]]) for op in rpush_list] + + async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]: + return [self.lists.pop(op["key"], None) for op in lpop_list] + + +def _postgres_rejection(sqlstate: str) -> RawQueryError: + return RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}} ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lock_held_by_another_pod", + [None, DB_SPEND_UPDATE_JOB_NAME, TOOL_DISCOVERY_COMMIT_JOB_NAME], + ids=["free", "spend_lock_held", "tool_lock_held"], +) +async def test_commit_with_redis_writes_every_workers_tools_in_one_statement_from_the_tool_lock_holder_only( + lock_held_by_another_pod: str | None, +): + redis = _InMemoryRedis() + if lock_held_by_another_pod is not None: + assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=lock_held_by_another_pod) + other_worker = ToolDiscoveryQueue() + other_worker.add_update({"tool_name": "get_weather", "origin": "mcp_server"}) + other_worker.add_update({"tool_name": "search", "origin": "user_defined"}) + await RedisUpdateBuffer(redis_cache=redis).store_tool_discoveries_in_redis(other_worker) + + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(return_value=2) + prisma_client.db.litellm_tooltable.upsert = AsyncMock() + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=True, + ) + + assert queue.flush() == [], "every worker hands its tools to redis each tick" + prisma_client.db.litellm_tooltable.upsert.assert_not_awaited() + if lock_held_by_another_pod == TOOL_DISCOVERY_COMMIT_JOB_NAME: + prisma_client.db.execute_raw.assert_not_awaited() + assert len(redis.lists[REDIS_TOOL_DISCOVERY_BUFFER_KEY]) == 2, "non-holders leave their tools for the holder" + return + assert REDIS_TOOL_DISCOVERY_BUFFER_KEY not in redis.lists, "the spend lock must not gate the tool drain" + prisma_client.db.execute_raw.assert_awaited_once() + sql, *params = prisma_client.db.execute_raw.call_args.args + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') and 'ON CONFLICT ("tool_name")' in sql + rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)} + assert rows == {"get_weather": 3, "search": 1}, f"one row per tool summed across workers, got {params}" + + +@pytest.mark.asyncio +async def test_commit_without_redis_writes_only_this_workers_tools_in_one_statement_and_never_touches_redis(): + redis = _InMemoryRedis() + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(return_value=2) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=False, + ) + + assert redis.lists == {}, "without the buffer flag nothing goes through redis" + assert queue.flush() == [] + prisma_client.db.execute_raw.assert_awaited_once() + sql, *params = prisma_client.db.execute_raw.call_args.args + assert 'ON CONFLICT ("tool_name")' in sql + rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)} + assert rows == {"get_weather": 2, "search": 1}, f"got {params}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "requeued"), + [ + (Exception("connection reset"), True), + (_postgres_rejection("23502"), False), + ], + ids=["connection_error_goes_back_to_redis", "rejected_row_is_dropped"], +) +async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably_did_not_apply( + failure: Exception, requeued: bool +): + redis = _InMemoryRedis() + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(side_effect=failure) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=True, + ) + + prisma_client.db.execute_raw.assert_awaited_once() + left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] + expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] + assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" diff --git a/tests/unit/proxy/db/test_tool_table_bulk_upsert.py b/tests/unit/proxy/db/test_tool_table_bulk_upsert.py new file mode 100644 index 00000000000..d1e79030f9d --- /dev/null +++ b/tests/unit/proxy/db/test_tool_table_bulk_upsert.py @@ -0,0 +1,38 @@ +"""Tests for the single-statement LiteLLM_ToolTable upsert (LIT-9997).""" + +from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert + + +def test_one_row_per_item_and_the_conflict_update_only_moves_the_counters(): + sql, params = build_tool_bulk_upsert( + [ + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}, + { + "tool_name": "tool_b", + "origin": "user_defined", + "created_by": "alice", + "call_count": 1, + "key_hash": "hash", + "team_id": "team", + "key_alias": "alias", + "user_agent": "agent", + }, + ] + ) + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') + assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql + assert "$4::integer" in sql and "$14::integer" in sql + assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql + assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql + assert '"last_used_at" = EXCLUDED."last_used_at"' in sql + assert '"input_policy"' not in sql and '"output_policy"' not in sql, "policies stay on the table defaults" + assert '"origin" =' not in sql and '"key_hash" =' not in sql, "an existing row keeps who discovered it" + assert len(params) == 20 + assert params[1:10] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None) + assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", "hash", "team", "alias", "agent") + assert params[0] != params[10], "each row gets its own tool_id" + + +def test_defaults_fill_origin_created_by_and_call_count(): + _, params = build_tool_bulk_upsert([{"tool_name": "tool_a"}]) + assert params[1:6] == ("tool_a", "user_defined", 1, "system", "system") diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index 021eb289ad7..4384507e841 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -126,6 +126,51 @@ def test_proxy_decisions_route_returns_answers_and_cost( assert upstream.calls[0].request.headers["authorization"] == "Bearer test-key" +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/systemone")) +@pytest.mark.parametrize("question_count", (128, 129, 154)) +def test_proxy_systemone_preserves_large_question_batches( + client: TestClient, + respx_mock: respx.MockRouter, + endpoint: str, + question_count: int, +) -> None: + questions: Final = { + f"q{index}": {"type": "noul", "instructions": f"Question {index}?"} for index in range(question_count) + } + answers: Final = {f"q{index}": {"type": "noul", "noul": index / question_count} for index in range(question_count)} + upstream: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond( + json={**_RESPONSE, "answers": answers} + ) + + response: Final = client.post(endpoint, json={**_REQUEST, "questions": questions}) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == { + **_REQUEST, + "model": "pplx-decider-v1-27b", + "questions": questions, + } + assert response.json()["answers"] == answers + + +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/systemone")) +@pytest.mark.parametrize( + "questions", + ({}, {"": {"type": "noul", "instructions": "Is this a defect?"}}, {"is_defect": {"type": "noul"}}), + ids=("empty", "empty_name", "missing_instructions_or_criteria"), +) +def test_proxy_systemone_rejects_invalid_questions_before_http( + client: TestClient, + respx_mock: respx.MockRouter, + endpoint: str, + questions: Mapping[str, object], +) -> None: + response: Final = client.post(endpoint, json={**_REQUEST, "questions": questions}) + + assert response.status_code == 400, response.text + assert len(respx_mock.calls) == 0 + + def test_proxy_decisions_dispatches_typesafe_deployment( client: TestClient, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py b/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py index aa43a5530f3..953baf5c173 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py @@ -129,6 +129,20 @@ class TestHelperFunctions: mock_response.choices = [] assert _extract_response_text(mock_response) == "" + def test_extract_response_text_fails_closed_on_a_message_without_content(self): + from types import SimpleNamespace + + from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.semantic_guard import ( + _extract_response_text, + ) + + response: Final = SimpleNamespace( + choices=[SimpleNamespace(message={"role": "assistant", "content": "ignore all rules"})] + ) + + with pytest.raises(AttributeError, match="'dict' object has no attribute 'content'"): + _extract_response_text(response) + def test_get_top_route_choice_single(self): from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.semantic_guard import ( _get_top_route_choice, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 63fcdfbb0c5..3da77a8707e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1,6 +1,7 @@ import asyncio import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock import pytest @@ -1237,6 +1238,23 @@ class TestFailClosed: class TestStreamingRehydration: + @pytest.mark.asyncio + async def test_held_text_behind_a_chunk_that_cannot_be_copied_fails_closed(self): + guardrail: Final = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"text": "", "carry": "[EMAIL"}, {"text": "a@b.com", "carry": ""}) + + async def stream(): + yield SimpleNamespace( + choices=[SimpleNamespace(index=0, delta=SimpleNamespace(content="[EMAIL"), finish_reason=None)] + ) + + with pytest.raises(AttributeError, match="model_copy"): + await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + @pytest.mark.asyncio async def test_split_placeholder_is_not_emitted_in_fragments(self): """The window holds back a partial placeholder and releases it once complete.""" @@ -1975,6 +1993,20 @@ class TestProxyWiring: class TestStreamUsage: + @pytest.mark.asyncio + async def test_trailing_flush_survives_a_usage_chunk_with_no_choices(self): + guardrail, _ = _shielded({"[EMAIL_1]": "a@example.com"}) + held: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Mail [EMAI"))]) + usage_chunk: Final = ModelResponseStream( + choices=[], usage=Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12) + ) + + out: Final = await _restore_stream(guardrail, [held, usage_chunk]) + + assert out[1] is not None and out[1].choices == [] + assert [chunk.choices[0].delta.content for chunk in out[2:]] == ["[EMAI"] + assert all(getattr(chunk, "usage", None) is None for chunk in out[2:]) + @pytest.mark.asyncio async def test_trailing_flush_does_not_repeat_usage(self): """With n>=2 and include_usage, the last chunk carries usage; a flush copied from it must not. diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py new file mode 100644 index 00000000000..9a7670dbe97 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py @@ -0,0 +1,285 @@ +""" +Tests for the Reco guardrail integration. +""" + +from collections.abc import Mapping +from typing import Final +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy.guardrails.guardrail_hooks.reco import ( + RecoGuardrail, + initialize_guardrail, +) +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.proxy.guardrails.guardrail_hooks.reco import ( + RecoConfigModel, + RecoOptionalParams, +) +from litellm.types.utils import GenericGuardrailAPIInputs + + +@pytest.fixture +def reco_guardrail() -> RecoGuardrail: + return RecoGuardrail( + reco_tenant_id="11111111-1111-1111-1111-111111111111", + api_base="https://edge1.us.reco.ai", + guardrail_name="my-reco-guardrail", + event_hook="pre_call", + default_on=True, + ) + + +class TestRecoGuardrailConfiguration: + def test_tenant_id_forwarded_as_header(self, reco_guardrail: RecoGuardrail) -> None: + assert reco_guardrail.headers == {"X-Reco-Tenant-Id": "11111111-1111-1111-1111-111111111111"} + + def test_tenant_header_merged_with_existing_headers(self) -> None: + guardrail: Final = RecoGuardrail( + reco_tenant_id="11111111-1111-1111-1111-111111111111", + api_base="https://edge1.us.reco.ai", + headers={"X-Custom": "value"}, + ) + assert guardrail.headers == { + "X-Custom": "value", + "X-Reco-Tenant-Id": "11111111-1111-1111-1111-111111111111", + } + + def test_missing_reco_tenant_id_raises(self) -> None: + with pytest.raises(ValueError, match="reco_tenant_id"): + RecoGuardrail(reco_tenant_id=None, api_base="https://edge1.us.reco.ai") + + def test_missing_api_base_raises(self) -> None: + with pytest.raises(ValueError, match="api_base"): + RecoGuardrail(reco_tenant_id="11111111-1111-1111-1111-111111111111", api_base=None) + + def test_non_uuid_reco_tenant_id_raises(self) -> None: + with pytest.raises(ValueError, match="UUID"): + RecoGuardrail(reco_tenant_id="tenant-123", api_base="https://edge1.us.reco.ai") + + def test_optional_params_rejects_non_uuid_tenant_id(self) -> None: + with pytest.raises(ValidationError, match="UUID"): + RecoOptionalParams(reco_tenant_id="tenant-123", api_base="https://edge1.us.reco.ai") + + def test_get_supported_event_hooks_returns_pre_call_only(self) -> None: + assert RecoGuardrail.get_supported_event_hooks() == [GuardrailEventHooks.pre_call] + + def test_get_config_model_returns_reco_config_model(self) -> None: + assert RecoGuardrail.get_config_model() is RecoConfigModel + + def test_config_model_ui_friendly_name(self) -> None: + assert RecoConfigModel.ui_friendly_name() == "Reco" + + def test_litellm_params_stay_assignable_for_registry_secret_resolution(self) -> None: + litellm_params: Final = LitellmParams(guardrail="reco", mode="pre_call") + + litellm_params.api_base = "https://edge1.us.reco.ai" + + assert litellm_params.api_base == "https://edge1.us.reco.ai" + + +class TestRecoGuardrailInitializer: + def test_initialize_guardrail_wires_optional_params(self) -> None: + litellm_params: Final = LitellmParams( + guardrail="reco", + mode="pre_call", + optional_params={ + "reco_tenant_id": "22222222-2222-2222-2222-222222222222", + "api_base": "https://edge2.eu.reco.ai", + }, + ) + + guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"}) + + assert isinstance(guardrail, RecoGuardrail) + assert guardrail.guardrail_name == "prod-reco" + assert guardrail.headers == {"X-Reco-Tenant-Id": "22222222-2222-2222-2222-222222222222"} + assert guardrail.api_base == "https://edge2.eu.reco.ai/beta/litellm_basic_guardrail_api" + + def test_initialize_guardrail_reads_plain_dict_optional_params(self) -> None: + """Guardrails reconciled from the DB can carry `optional_params` as a + plain dict rather than a pydantic object; attribute-only access would + silently read None for every field.""" + litellm_params: Final = LitellmParams.model_construct( + guardrail="reco", + optional_params={ + "reco_tenant_id": "33333333-3333-3333-3333-333333333333", + "api_base": "https://edge3.eu.reco.ai", + }, + mode="pre_call", + default_on=True, + ) + + guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"}) + + assert guardrail.headers == {"X-Reco-Tenant-Id": "33333333-3333-3333-3333-333333333333"} + assert guardrail.api_base == "https://edge3.eu.reco.ai/beta/litellm_basic_guardrail_api" + + def test_initialize_guardrail_resolves_env_var_reference(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("RECO_TEST_TENANT_ID", "44444444-4444-4444-4444-444444444444") + litellm_params: Final = LitellmParams( + guardrail="reco", + optional_params={ + "reco_tenant_id": "os.environ/RECO_TEST_TENANT_ID", + "api_base": "https://edge4.eu.reco.ai", + }, + mode="pre_call", + default_on=True, + ) + + guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"}) + + assert guardrail.headers == {"X-Reco-Tenant-Id": "44444444-4444-4444-4444-444444444444"} + + +def _litellm_params(top_level: Mapping[str, object], nested: Mapping[str, object]) -> LitellmParams: + return LitellmParams( + guardrail="reco", + mode="pre_call", + optional_params={ + "reco_tenant_id": "22222222-2222-2222-2222-222222222222", + "api_base": "https://edge2.eu.reco.ai", + **nested, + }, + **top_level, + ) + + +async def _apply_with_reco_unreachable(guardrail: RecoGuardrail) -> GenericGuardrailAPIInputs: + with patch.object(guardrail.async_handler, "post", side_effect=httpx.ConnectError("connection refused")): + return await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {}}, + input_type="request", + ) + + +class TestRecoGuardrailFailureModes: + @pytest.mark.asyncio + async def test_unreachable_reco_blocks_request_by_default(self, reco_guardrail: RecoGuardrail) -> None: + with pytest.raises(Exception, match="connection refused"): + await _apply_with_reco_unreachable(reco_guardrail) + + @pytest.mark.asyncio + async def test_fail_open_lets_request_through_when_reco_unreachable(self) -> None: + guardrail: Final = RecoGuardrail( + reco_tenant_id="11111111-1111-1111-1111-111111111111", + api_base="https://edge1.us.reco.ai", + fail_on_error=False, + ) + + result: Final = await _apply_with_reco_unreachable(guardrail) + + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_fail_closed_blocks_request_when_reco_unreachable(self) -> None: + guardrail: Final = RecoGuardrail( + reco_tenant_id="11111111-1111-1111-1111-111111111111", + api_base="https://edge1.us.reco.ai", + unreachable_fallback="fail_closed", + fail_on_error=True, + ) + + with pytest.raises(Exception, match="connection refused"): + await _apply_with_reco_unreachable(guardrail) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("top_level", "nested"), + [ + pytest.param({}, {}, id="defaults"), + pytest.param({}, {"unreachable_fallback": "fail_closed", "fail_on_error": True}, id="optional-params"), + pytest.param({"unreachable_fallback": "fail_closed", "fail_on_error": True}, {}, id="top-level"), + ], + ) + async def test_initialize_guardrail_fails_closed_when_configured( + self, top_level: Mapping[str, object], nested: Mapping[str, object] + ) -> None: + guardrail: Final = initialize_guardrail(_litellm_params(top_level, nested), {"guardrail_name": "prod-reco"}) + + with pytest.raises(Exception, match="connection refused"): + await _apply_with_reco_unreachable(guardrail) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("top_level", "nested"), + [ + pytest.param({}, {"fail_on_error": False}, id="optional-params"), + pytest.param({"fail_on_error": False}, {}, id="top-level"), + pytest.param( + {"unreachable_fallback": "fail_closed", "fail_on_error": True}, + {"fail_on_error": False}, + id="optional-params-wins-over-top-level", + ), + ], + ) + async def test_initialize_guardrail_fails_open_when_configured( + self, top_level: Mapping[str, object], nested: Mapping[str, object] + ) -> None: + guardrail: Final = initialize_guardrail(_litellm_params(top_level, nested), {"guardrail_name": "prod-reco"}) + + result: Final = await _apply_with_reco_unreachable(guardrail) + + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_initialize_guardrail_keeps_top_level_fail_open_when_nested_params_are_a_model(self) -> None: + nested: Final = RecoOptionalParams( + reco_tenant_id="22222222-2222-2222-2222-222222222222", api_base="https://edge2.eu.reco.ai" + ) + litellm_params: Final = LitellmParams.model_construct( + guardrail="reco", mode="pre_call", optional_params=nested, fail_on_error=False + ) + + guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"}) + result: Final = await _apply_with_reco_unreachable(guardrail) + + assert result["texts"] == ["hello"] + + +class TestRecoGuardrailBlocking: + @pytest.mark.asyncio + async def test_action_blocked_raises_exception_with_configured_guardrail_name( + self, reco_guardrail: RecoGuardrail + ) -> None: + mock_response: Final = MagicMock() + mock_response.json.return_value = { + "action": "BLOCKED", + "blocked_reason": "Sensitive data detected", + } + mock_response.raise_for_status = MagicMock() + + with patch.object(reco_guardrail.async_handler, "post", return_value=mock_response) as mock_post: + with pytest.raises(GuardrailRaisedException) as exc_info: + await reco_guardrail.apply_guardrail( + inputs={"texts": ["some sensitive prompt"]}, + request_data={"metadata": {}}, + input_type="request", + ) + + assert exc_info.value.guardrail_name == "my-reco-guardrail" + assert str(exc_info.value) == "Sensitive data detected" + assert exc_info.value.blocked_content is True + + _, call_kwargs = mock_post.call_args + assert call_kwargs["headers"]["X-Reco-Tenant-Id"] == "11111111-1111-1111-1111-111111111111" + + @pytest.mark.asyncio + async def test_action_none_allows_content(self, reco_guardrail: RecoGuardrail) -> None: + mock_response: Final = MagicMock() + mock_response.json.return_value = {"action": "NONE", "texts": ["hello"]} + mock_response.raise_for_status = MagicMock() + + with patch.object(reco_guardrail.async_handler, "post", return_value=mock_response): + result: Final = await reco_guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {}}, + input_type="request", + ) + + assert result["texts"] == ["hello"] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 3efbe6475b2..a132ce8eada 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3,6 +3,7 @@ import asyncio import contextlib import io +import json import logging from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator from types import SimpleNamespace @@ -35,6 +36,7 @@ from litellm.llms.openai.chat.guardrail_translation.handler import ( from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) +from litellm.llms.pass_through.guardrail_translation.handler import PassThroughEndpointHandler from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) @@ -3053,3 +3055,34 @@ async def test_pre_call_hook_skips_attachment_scan_when_guardrail_has_none(): ) assert len(guardrail.apply_calls) == 1 + + +@pytest.mark.parametrize("hook", ["async_pre_call_hook", "async_moderation_hook"]) +@pytest.mark.asyncio +async def test_pass_through_guardrail_text_is_the_client_body_without_the_injected_key_metadata( + monkeypatch: pytest.MonkeyPatch, hook: str +) -> None: + _patch_translation_mappings(monkeypatch, {CallTypes.pass_through: PassThroughEndpointHandler}) + guardrail: Final = RecordingGuardrail() + body: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Say hi in one word"}], + "guardrails": ["recording-guardrail"], + } + data: Final = {**body, "guardrail_to_apply": guardrail} + user_api_key_dict: Final = UserAPIKeyAuth( + token="hashed-key", user_email="someone@example.com", team_id="team-1", user_role="proxy_admin" + ) + + if hook == "async_pre_call_hook": + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=CallTypes.pass_through.value + ) + else: + await UnifiedLLMGuardrails().async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=CallTypes.pass_through.value + ) + + texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls] + assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts + assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"] diff --git a/tests/unit/proxy/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py index 920ffc77095..c14375135dc 100644 --- a/tests/unit/proxy/guardrails/test_content_utils.py +++ b/tests/unit/proxy/guardrails/test_content_utils.py @@ -1,5 +1,7 @@ """Tests for the shared guardrail content extraction helpers.""" +import json + from litellm.proxy.guardrails._content_utils import ( apply_redacted_messages_back, build_inspection_messages, @@ -167,6 +169,21 @@ def test_iter_message_text_inspects_reasoning_content_and_summary(): assert list(iter_message_text(data)) == ["content secret", "summary secret"] + +def test_iter_message_text_inspects_compaction_summary(): + data = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "banned word"}), + }, + {"role": "user", "content": "next turn"}, + ] + } + assert list(iter_message_text(data)) == ["banned word", "next turn"] + + # ── walk_user_text ──────────────────────────────────────────────────────────── @@ -347,6 +364,30 @@ def test_walk_user_text_redacts_reasoning_content_and_summary(): assert item["id"] == "rs_1" + +def test_walk_user_text_redacts_compaction_summary_and_keeps_the_block(): + data = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps( + {"type": "compaction", "content": "AKIAEXAMPLE summary", "signature": "sig"} + ), + } + ] + } + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + item = data["input"][0] + assert visited == 1 + assert (item["type"], item["id"]) == ("compaction", "cmp_1") + assert json.loads(item["encrypted_content"]) == { + "type": "compaction", + "content": "[REDACTED] summary", + "signature": "sig", + } + + # ── build_inspection_messages ───────────────────────────────────────────────── diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 950d1742d49..67b280eb379 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -5437,3 +5437,64 @@ async def test_health_services_endpoint_langfuse_missing_keys_errors(monkeypatch with pytest.raises(ProxyException, match="auth_check failed") as raised: await health_services_endpoint(service="langfuse") assert "connection refused by lf.internal.example" in str(raised.value.message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("admin", [False, True]) +async def test_unsaved_chatgpt_account_probe_requires_proxy_admin( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, admin: bool +) -> None: + from starlette.requests import Request + from litellm.proxy import proxy_server + + monkeypatch.setenv("CHATGPT_AUTH_FILE", str(tmp_path / "missing-default-auth.json")) + auth_file: Final = tmp_path / "auth.json" + auth_file.write_text( + json.dumps({"access_token": "token-probe", "account_id": "acct-probe", "expires_at": 4102444800}) + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + output: Final = { + "id": "msg_probe", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + response: Final = { + "id": "resp_probe", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [output], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + with respx.mock(assert_all_called=False) as upstream: + upstream.post("https://chatgpt.com/backend-api/codex/responses").respond( + 200, + text="data: " + json.dumps({"type": "response.completed", "response": response}) + "\n\n", + headers={"content-type": "text/event-stream"}, + ) + call: Final = health_test_model_connection( + request=Request({"type": "http", "headers": []}), + mode="responses", + litellm_params={"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(auth_file)}, + model_info={}, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN if admin else LitellmUserRoles.INTERNAL_USER, token="test-key" + ), + ) + if admin: + result: Final = await call + assert result["status"] == "success" + assert upstream.calls.last.request.headers["Authorization"] == "Bearer token-probe" + assert upstream.calls.last.request.headers["ChatGPT-Account-Id"] == "acct-probe" + assert isinstance(json.loads(upstream.calls.last.request.content)["input"], list) + else: + with pytest.raises(ProxyException) as error: + await call + assert error.value.code == "403" + assert len(upstream.calls) == 0 diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 8a74f2d50b4..345f37f1db4 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3,10 +3,11 @@ Unit Tests for the max parallel request limiter v3 for the proxy """ import asyncio +import contextvars import logging import sys import time -from collections.abc import Iterator, Sequence +from collections.abc import AsyncGenerator, Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone from typing import Any, Dict, Final, List, Optional @@ -54,6 +55,7 @@ from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from litellm.types.utils import ( + CallTypes, EmbeddingResponse, ModelResponse, PromptTokensDetailsWrapper, @@ -4206,6 +4208,175 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: ) == 0 +@pytest.mark.asyncio +async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() -> None: + api_key: Final = hash_token("sk-mcp-auto-execute") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = "mcp-auto-execute-owner" + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + def mcp_tool_call() -> dict[str, object]: + return { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Tool: add"}], + "mcp_server_name": "math", + } + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + async def pre_call_from_another_request( + data: dict[str, object], call_type: str + ) -> None: + await contextvars.Context().run(asyncio.create_task, pre_call(data, call_type)) + + await pre_call( + {"model": "gpt-4o-mini", "input": "add 2 and 3", "litellm_call_id": owner_call_id}, + CallTypes.aresponses.value, + ) + await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value) + await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value) + assert await in_flight() == 1 + + with pytest.raises(HTTPException): + await pre_call_from_another_request( + mcp_tool_call(), CallTypes.call_mcp_tool.value + ) + + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id}, + response_obj=None, + start_time=None, + end_time=None, + ) + assert await in_flight() == 0 + await pre_call_from_another_request( + {"model": "gpt-4o-mini", "input": "next request"}, + CallTypes.aresponses.value, + ) + assert await in_flight() == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stream", "in_flight_after_first_round_logs"), + [(False, 1), (True, 0)], + ids=["non_stream_round_logged_late_keeps_slot", "stream_logged_at_its_end_releases_slot"], +) +async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request_end( + stream: bool, in_flight_after_first_round_logs: int +) -> None: + reused_at: Final = datetime(2026, 1, 5, 12, 0, 0) + api_key: Final = hash_token(f"sk-mcp-late-round-log-{stream}") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=lambda: reused_at, + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = f"mcp-late-round-log-{stream}" + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + async def round_logged(ended_at: datetime) -> None: + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id, "stream": stream}, + response_obj=None, + start_time=None, + end_time=ended_at, + ) + + await pre_call( + { + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "add 2 and 3"}], + "litellm_call_id": owner_call_id, + }, + CallTypes.anthropic_messages.value, + ) + await pre_call( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ) + + await round_logged(reused_at - timedelta(seconds=1)) + assert await in_flight() == in_flight_after_first_round_logs + + await round_logged(reused_at + timedelta(seconds=1)) + assert await in_flight() == 0 + + +@pytest.mark.asyncio +async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None: + api_key: Final = hash_token("sk-mcp-reuse-headers") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) + data: Final[dict[str, object]] = { + "model": "gpt-4o-mini", + "input": "add 2 and 3", + "litellm_call_id": "mcp-reuse-headers-owner", + } + for request_data, call_type in ( + (data, CallTypes.aresponses.value), + ( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ), + ): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=request_data, + call_type=call_type, + ) + response: Final = ModelResponse( + id="mcp-reuse-headers", + object="chat.completion", + created=0, + model="gpt-4o-mini", + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + choices=[], + ) + response._hidden_params = {} + + await handler.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + headers: Final = response._hidden_params["additional_headers"] + assert headers["x-ratelimit-api_key-limit-max_parallel_requests"] == 2 + assert headers["x-ratelimit-api_key-remaining-max_parallel_requests"] == 1 + + @pytest.mark.asyncio async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None: from unittest.mock import AsyncMock @@ -4473,6 +4644,140 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): ) +_SSE_KEEPALIVE_CALLBACK_TIME: Final = datetime(2026, 1, 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("after_ping", [False, True], ids=["fast", "pinged"]) +@pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"]) +@pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"]) +async def test_sse_keepalive_releases_parallel_capacity( + after_ping: bool, fails: bool, preadmitted: bool +) -> None: + from fastapi.responses import StreamingResponse + + from litellm.proxy.common_request_processing import open_sse_before_first_byte + from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES + + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + api_key: Final = hash_token("sk-keepalive-release") + auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + counter: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + + async def request(call_id: str) -> None: + data: Final = {"model": "test-model", "litellm_call_id": call_id} + upstream_answers: Final = asyncio.Event() + if not after_ping: + upstream_answers.set() + + async def admit() -> None: + await handler.async_pre_call_hook( + user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion" + ) + + async def failure(exc: Exception) -> None: + await handler.async_post_call_failure_hook( + request_data=data, original_exception=exc, user_api_key_dict=auth + ) + + async def stream() -> AsyncGenerator[bytes, None]: + yield b"data: [DONE]\n\n" + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": call_id, + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + }, + response_obj=ModelResponse(usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)), + start_time=_SSE_KEEPALIVE_CALLBACK_TIME, + end_time=_SSE_KEEPALIVE_CALLBACK_TIME, + ) + + async def produce() -> StreamingResponse: + if not preadmitted: + await admit() + await upstream_answers.wait() + if fails: + raise HTTPException(status_code=502, detail="upstream failed") + return StreamingResponse(stream()) + + if preadmitted: + await admit() + response_call: Final = open_sse_before_first_byte( + produce(), 0.001 if after_ping else 60.0, on_late_failure=failure + ) + if fails and not after_ping: + with pytest.raises(HTTPException) as rejected: + await response_call + assert rejected.value.status_code == 502 + await failure(rejected.value) + return + response: Final = await response_call + first: Final = await anext(response.body_iterator) + if after_ping: + upstream_answers.set() + rest: Final = tuple([chunk async for chunk in response.body_iterator]) + assert first == (SSE_COMMENT_PING_BYTES if after_ping else b"data: [DONE]\n\n") + if fails: + error_frame: Final = rest[-2] + assert isinstance(error_frame, bytes) and b"upstream failed" in error_frame + assert rest[-1] == b"data: [DONE]\n\n" + elif after_ping: + assert rest[-1] == b"data: [DONE]\n\n" + else: + assert rest == () + + for index in range(3): + await contextvars.Context().run(asyncio.create_task, request(f"keepalive-{index}")) + assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0 + + +@pytest.mark.asyncio +async def test_sse_keepalive_rejects_overlap_and_reuses_completed_capacity() -> None: + from fastapi.responses import StreamingResponse + + from litellm.proxy.common_request_processing import open_sse_before_first_byte + + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + api_key: Final = hash_token("sk-keepalive-overlap") + auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + admitted: Final = asyncio.Event() + finish: Final = asyncio.Event() + + async def request(call_id: str) -> None: + async def stream() -> AsyncGenerator[bytes, None]: + await finish.wait() + yield b"data: [DONE]\n\n" + await handler.async_post_call_success_hook( + data={"litellm_call_id": call_id}, user_api_key_dict=auth, response=ModelResponse() + ) + + async def produce() -> StreamingResponse: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": "test-model", "litellm_call_id": call_id}, + call_type="acompletion", + ) + admitted.set() + return StreamingResponse(stream()) + + response: Final = await open_sse_before_first_byte(produce(), 60.0) + assert tuple([chunk async for chunk in response.body_iterator]) == (b"data: [DONE]\n\n",) + + first: Final = contextvars.Context().run(asyncio.create_task, request("first")) + await admitted.wait() + try: + with pytest.raises(HTTPException) as rejected: + await contextvars.Context().run(asyncio.create_task, request("overlapping")) + assert rejected.value.status_code == 429 + finally: + finish.set() + await first + await contextvars.Context().run(asyncio.create_task, request("after-completion")) + + @pytest.mark.asyncio async def test_read_only_gauge_check_counts_without_acquiring_v3(): """ diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 36072d7eb79..40857f7c419 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,36 +1,44 @@ import asyncio import json import logging +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, create_autospec, patch import httpx import pytest +from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY -from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth +from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( - _get_budget_reservation_from_metadata, ProxyDBLogger, + _get_budget_reservation_from_metadata, _should_track_cost_callback, _update_database_and_spend_counters, + is_pre_routing_rejection, run_spend_event, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError -from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.utils import ( + REJECTED_BEFORE_REQUEST_SETUP, + REJECTED_BEFORE_REQUEST_SETUP_KEY, + BilledFailedAttempt, + ProxyUpdateSpend, +) from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage @@ -416,6 +424,74 @@ async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_m assert metadata["standard_logging_guardrail_information"] == metadata_bucket_info +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_headers, caller_metadata, redaction_is_global, code_word_is_stored", + [ + ({"x-litellm-enable-message-redaction": "true"}, {"user_id": "claude-code-session"}, False, False), + ({"litellm-enable-message-redaction": "true"}, {}, False, False), + ({"user-agent": "claude-cli"}, {"user_id": "claude-code-session"}, False, True), + ({"user-agent": "claude-cli"}, {"headers": {"x-litellm-enable-message-redaction": "true"}}, False, False), + ({"litellm-disable-message-redaction": "true"}, {"user_id": "claude-code-session"}, True, False), + ], + ids=[ + "enable_header_redacts", + "legacy_enable_header_redacts", + "no_header_stores_prompt", + "caller_metadata_headers_still_apply", + "disable_header_never_loosens_global_redaction", + ], +) +async def test_async_post_call_failure_hook_honors_the_redaction_header_on_a_litellm_metadata_route( + monkeypatch: pytest.MonkeyPatch, + request_headers: dict[str, str], + caller_metadata: dict[str, object], + redaction_is_global: bool, + code_word_is_stored: bool, +): + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + monkeypatch.setattr(litellm, "turn_off_message_logging", redaction_is_global) + writer: Final = MagicMock(spec=DBSpendUpdateWriter) + writer.update_database = AsyncMock() + logger: Final = ProxyDBLogger(spend_writer=lambda: writer) + request_input: Final = "The code word is VERMILION-42. Reply OK." + body_metadata: Final = dict(caller_metadata) + litellm_metadata_headers: Final = dict(request_headers) + + await logger.async_post_call_failure_hook( + request_data={ + "model": "gpt-5.6", + "input": request_input, + "metadata": body_metadata, + "litellm_metadata": {"headers": litellm_metadata_headers}, + "proxy_server_request": { + "url": "http://localhost/v1/responses", + "body": { + "model": "gpt-5.6", + "input": request_input, + "temperature": 7, + "metadata": body_metadata, + "litellm_metadata": {"headers": litellm_metadata_headers}, + }, + }, + }, + original_exception=Exception("temperature must be at most 2"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + row: Final = get_logging_payload( + kwargs=writer.update_database.call_args.kwargs["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + stored_request: Final = json.loads(row["proxy_server_request"]) + assert stored_request["temperature"] == 7 + assert ("VERMILION-42" in row["proxy_server_request"]) is code_word_is_stored + assert stored_request["metadata"].get("headers") == caller_metadata.get("headers") + assert stored_request["litellm_metadata"]["headers"] == request_headers + + @pytest.mark.asyncio @pytest.mark.parametrize( "used_client_oauth_token, custom_llm_provider, expected", @@ -3724,3 +3800,103 @@ async def test_batch_accounting_does_not_retry_after_a_search_allowance_update(m assert sum(counter.totals.values()) == 5 assert outcome.failed is True assert outcome.retryable_failure is False + + +_ROUTED_REQUEST: Final[Mapping[str, object]] = { + "model": "anthropic-haiku-4-5", + "messages": [{"role": "user", "content": "hi"}], + "proxy_server_request": {"request_id": "req-1"}, + "litellm_call_id": "call-1", +} +_PRE_ROUTING_REQUEST: Final[Mapping[str, object]] = { + **_ROUTED_REQUEST, + REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP, +} +_SPOOFED_REQUEST: Final[Mapping[str, object]] = {**_ROUTED_REQUEST, REJECTED_BEFORE_REQUEST_SETUP_KEY: True} +_INVALID_REQUEST_400: Final = ProxyException( + message="Violated guardrail policy", type="invalid_request_error", param=None, code=400 +) +_PROVIDER_AUTH_ERROR: Final = litellm.AuthenticationError( + message="bad key", llm_provider="anthropic", model="claude-haiku-4-5" +) + + +@pytest.mark.parametrize( + ("exception", "request_data", "expected"), + [ + (_INVALID_REQUEST_400, _PRE_ROUTING_REQUEST, True), + (litellm.BudgetExceededError(current_cost=1.0, max_budget=0.0), _PRE_ROUTING_REQUEST, True), + (ProxyModelNotFoundError(route="/chat/completions", model_name="4.6"), _ROUTED_REQUEST, True), + (_INVALID_REQUEST_400, _ROUTED_REQUEST, False), + (_INVALID_REQUEST_400, _SPOOFED_REQUEST, False), + (_PROVIDER_AUTH_ERROR, _ROUTED_REQUEST, False), + ( + litellm.RateLimitError(message="slow down", llm_provider="anthropic", model="claude-haiku-4-5"), + _ROUTED_REQUEST, + False, + ), + (HTTPException(status_code=429, detail="Max parallel request limit reached"), _ROUTED_REQUEST, False), + ], + ids=[ + "bad-body", + "budget", + "unknown-model", + "guardrail-block", + "guardrail-block-with-client-spoofed-marker", + "provider-auth", + "provider-429", + "proxy-429", + ], +) +def test_is_pre_routing_rejection_is_decided_by_pipeline_stage_not_exception_type( + exception: Exception, request_data: Mapping[str, object], expected: bool +) -> None: + assert is_pre_routing_rejection(exception, request_data) is expected + + +class _RecordingSpendWriter: + def __init__(self) -> None: + self.rows: list[Mapping[str, object]] = [] # mutable-ok: test double records the spend rows written + + async def update_database(self, **kwargs: object) -> None: + self.rows.append(kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("flag_enabled", "exception", "request_data", "rows_written"), + [ + (True, _INVALID_REQUEST_400, _PRE_ROUTING_REQUEST, 0), + (True, litellm.BudgetExceededError(current_cost=1.0, max_budget=0.0), _PRE_ROUTING_REQUEST, 0), + (True, ProxyModelNotFoundError(route="/chat/completions", model_name="4.6"), _ROUTED_REQUEST, 0), + (True, _INVALID_REQUEST_400, _ROUTED_REQUEST, 1), + (True, _INVALID_REQUEST_400, _SPOOFED_REQUEST, 1), + (True, _PROVIDER_AUTH_ERROR, _ROUTED_REQUEST, 1), + (False, _INVALID_REQUEST_400, _PRE_ROUTING_REQUEST, 1), + (False, litellm.BudgetExceededError(current_cost=1.0, max_budget=0.0), _PRE_ROUTING_REQUEST, 1), + ], + ids=[ + "flag-bad-body", + "flag-budget", + "flag-unknown-model", + "flag-guardrail-block-kept", + "flag-guardrail-block-with-spoofed-marker-kept", + "flag-provider-error-kept", + "off-bad-body", + "off-budget", + ], +) +async def test_async_post_call_failure_hook_disable_pre_routing_error_logs( + flag_enabled: bool, exception: Exception, request_data: Mapping[str, object], rows_written: int +) -> None: + writer: Final = _RecordingSpendWriter() + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="test_api_key", user_id="u", team_id="t", request_route="/chat/completions" + ) + with patch("litellm.proxy.proxy_server.general_settings", {"disable_pre_routing_error_logs": flag_enabled}): + await ProxyDBLogger(spend_writer=lambda: cast(DBSpendUpdateWriter, writer)).async_post_call_failure_hook( + request_data={**request_data, "metadata": dict[str, object]()}, + original_exception=exception, + user_api_key_dict=user_api_key_dict, + ) + assert len(writer.rows) == rows_written diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index bb2a89b1d98..3d1e494e36c 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -15,8 +15,9 @@ Redis. """ import asyncio +import json from datetime import datetime, timedelta -from typing import Any, Dict +from typing import Any, Dict, Final import pytest @@ -3727,3 +3728,50 @@ async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loo if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) + + +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_reservation_skips_the_history_a_compaction_item_replaces(rate_limiter): + handler, _cache = rate_limiter + compacted_request: Final = {"model": "gpt-4o", "input": _COMPACTED_INPUT, "max_output_tokens": 1} + + assert handler._estimate_input_and_output_tokens( + data=compacted_request, call_type="aresponses" + ) == handler._estimate_input_and_output_tokens( + data={**compacted_request, "input": _COMPACTED_INPUT[2:]}, call_type="aresponses" + ) + assert handler._estimate_precise_input_tokens( + compacted_request, model="gpt-4o", call_type="aresponses" + ) == handler._estimate_precise_input_tokens( + {**compacted_request, "input": _SUMMARY_INPUT}, model="gpt-4o", call_type="aresponses" + ) + + +def test_reservation_counts_a_compaction_item_by_its_summary(rate_limiter): + handler, _cache = rate_limiter + padded_item: Final = { + **_COMPACTED_INPUT[2], + "encrypted_content": json.dumps( + {"type": "compaction", "content": _COMPACTION_SUMMARY, "opaque_state": "x" * 4000} + ), + } + + assert handler._estimate_input_and_output_tokens( + data={"model": "gpt-4o", "input": [padded_item, _NEXT_TURN], "max_output_tokens": 1}, call_type="aresponses" + ) == handler._estimate_input_and_output_tokens( + data={"model": "gpt-4o", "input": _COMPACTED_INPUT[2:], "max_output_tokens": 1}, call_type="aresponses" + ) diff --git a/tests/unit/proxy/lens/test_buffering.py b/tests/unit/proxy/lens/test_buffering.py index b544b70d36b..2b868e72096 100644 --- a/tests/unit/proxy/lens/test_buffering.py +++ b/tests/unit/proxy/lens/test_buffering.py @@ -249,7 +249,10 @@ async def test_admission_releases_upload_on_disconnect_or_cancellation(cancel: b @pytest.mark.asyncio @pytest.mark.parametrize("reject_auth", [False, True]) -async def test_admitted_bytes_survive_auth_parsing_and_remain_reserved_until_sent(reject_auth: bool) -> None: +@pytest.mark.parametrize("body", [b" {} ", b"[]", b"null", b'"s"']) +async def test_admitted_bytes_survive_auth_parsing_and_remain_reserved_until_sent( + reject_auth: bool, body: bytes +) -> None: app: Final = FastAPI() router: Final = APIRouter(route_class=LensRoute) @@ -264,9 +267,11 @@ async def test_admitted_bytes_survive_auth_parsing_and_remain_reserved_until_sen if reject_auth: raise HTTPException(401, "Invalid key") - async with httpx.AsyncClient( - transport=httpx.MockTransport(lambda _: httpx.Response(200, content=b"result")) - ) as upstream: + def upstream_response(incoming: httpx.Request) -> httpx.Response: + assert incoming.content == body + return httpx.Response(200, content=b"result") + + async with httpx.AsyncClient(transport=httpx.MockTransport(upstream_response)) as upstream: @router.post("/lens/datasets") async def endpoint(request: Request, auth: Annotated[None, Depends(auth_reader)]) -> Response: @@ -284,14 +289,14 @@ async def test_admitted_bytes_survive_auth_parsing_and_remain_reserved_until_sen sending.set() await finish.wait() - incoming: Final = request(b" {} ") + incoming: Final = request(body) task: Final = asyncio.create_task(app(incoming.scope, incoming.receive, send)) try: await asyncio.wait_for(sending.wait(), 2) async with httpx.AsyncClient( transport=httpx.ASGITransport(app=app), base_url="http://gateway.test" ) as client: - blocked: Final = await client.post("/lens/datasets", content=b" {} " * 2) + blocked: Final = await client.post("/lens/datasets", content=b" " * (11 - len(body))) unrelated: Final = await client.post("/ordinary", content=b"ordinary body") assert blocked.status_code == 503 assert unrelated.content == b"ordinary body" diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_budgets.py b/tests/unit/proxy/management_endpoints/management_v1/test_budgets.py index 2b438a9d370..4b42e6404a3 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_budgets.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_budgets.py @@ -52,7 +52,7 @@ app.include_router(router) client = TestClient(app) BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets" -SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpd_limit", "tpm_limit"] +SORTABLE = ["budget_duration", "budget_id", "created_at", "max_budget", "rpm_limit", "tpd_limit", "tpm_limit"] def _row(budget_id: str, **overrides: Any) -> dict[str, Any]: @@ -172,22 +172,12 @@ def test_appends_the_tiebreaker_to_an_explicit_sort(query_raw, as_proxy_admin): _get("sort=-max_budget") - assert 'ORDER BY "max_budget" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0] + assert 'ORDER BY "max_budget" DESC NULLS FIRST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0] -def test_refuses_to_sort_on_budget_duration(query_raw, as_proxy_admin): - """The column holds "7d"/"30d", so a lexicographic ORDER BY would put "30d" - before "7d" and silently mis-order the page.""" - _serve(query_raw, []) - +def test_accepts_duration_sorting(query_raw, as_proxy_admin): response = _get("sort=budget_duration") - - assert response.status_code == 400 - assert response.headers["content-type"].startswith("application/problem+json") - body = response.json() - assert "budget_duration" in body["detail"] - assert body["allowed"] == SORTABLE - query_raw.assert_not_called() + assert response.status_code == 200 def test_the_advertised_sort_fields_are_the_ones_that_work(query_raw, as_proxy_admin): @@ -346,7 +336,7 @@ def test_q_does_not_search_any_other_column(query_raw, as_proxy_admin): def test_is_null_selects_the_unlimited_budgets(query_raw, as_proxy_admin): - """"Unlimited" is max_budget IS NULL; `max_budget = 0` would be a hard zero cap.""" + """ "Unlimited" is max_budget IS NULL; `max_budget = 0` would be a hard zero cap.""" _serve(query_raw, [_row("b-unlimited", max_budget=None)]) body = _get("filter[max_budget][is_null]=true").json() diff --git a/tests/unit/proxy/management_endpoints/policy_endpoints/test_endpoints.py b/tests/unit/proxy/management_endpoints/policy_endpoints/test_endpoints.py index 62863163b8e..4b62cf79094 100644 --- a/tests/unit/proxy/management_endpoints/policy_endpoints/test_endpoints.py +++ b/tests/unit/proxy/management_endpoints/policy_endpoints/test_endpoints.py @@ -5,14 +5,20 @@ Tests _test_guardrail_definitions and _compute_overall_action directly without needing a running proxy. """ +import json +from pathlib import Path +from unittest.mock import patch + import pytest from fastapi import Request +import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.management_endpoints.policy_endpoints.endpoints import ( GuardrailTestResultEntry, _compute_overall_action, _test_guardrail_definitions, + get_policy_templates, list_policies, ) from litellm.proxy.policy_engine.policy_registry import get_policy_registry @@ -368,3 +374,19 @@ async def test_list_policies_reports_each_loaded_policy_with_its_resolved_guardr assert response == PolicyListResponse(policies=expected_policies, total_count=0) assert list(response.policies) == list(expected_policies) + + +@pytest.mark.parametrize("flag", ["LITELLM_OFFLINE_MODE", "LITELLM_LOCAL_POLICY_TEMPLATES"]) +@pytest.mark.asyncio +async def test_policy_templates_serve_the_bundled_copy_without_fetching(monkeypatch: pytest.MonkeyPatch, flag: str): + monkeypatch.setenv(flag, "true") + bundled = json.loads((Path(litellm.__file__).parent / "policy_templates_backup.json").read_text(encoding="utf-8")) + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + side_effect=AssertionError("policy templates were fetched"), + ) as get_client: + response = await get_policy_templates(request=_policy_list_request(), user_api_key_dict=UserAPIKeyAuth()) + + get_client.assert_not_called() + assert response == bundled diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 303b8e060ed..a718426eea6 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -8933,7 +8933,14 @@ class TestNonAdminCannotPersistWifFieldsOnModel: mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() @pytest.mark.asyncio - async def test_patch_model_non_admin_cannot_set_wif_field(self): + @pytest.mark.parametrize( + "field,value", + [ + ("anthropic_keycloak_token_url", "https://identity.example.com/token"), + ("chatgpt_auth_file", "/tokens/account-b/auth.json"), + ], + ) + async def test_patch_model_non_admin_cannot_set_wif_field(self, field: str, value: str) -> None: from litellm.proxy.management_endpoints.model_management_endpoints import ( patch_model, ) @@ -8979,13 +8986,13 @@ class TestNonAdminCannotPersistWifFieldsOnModel: model_id="m1", patch_data=updateDeployment( litellm_params=updateLiteLLMParams( - anthropic_keycloak_token_url="https://attacker.example/token", + **{field: value}, ) ), user_api_key_dict=non_admin, ) err = exc_info.value - assert getattr(err, "param", "") == "anthropic_keycloak_token_url" + assert getattr(err, "param", "") == field mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() @pytest.mark.asyncio @@ -9044,7 +9051,14 @@ class TestNonAdminCannotPersistWifFieldsOnModel: mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() @pytest.mark.asyncio - async def test_patch_model_admin_can_set_wif_field(self): + @pytest.mark.parametrize( + "field,value", + [ + ("anthropic_keycloak_token_url", "https://identity.example.com/token"), + ("chatgpt_auth_file", "/tokens/account-b/auth.json"), + ], + ) + async def test_patch_model_admin_can_set_wif_field(self, field: str, value: str) -> None: from litellm.proxy.management_endpoints.model_management_endpoints import ( patch_model, ) @@ -9095,7 +9109,7 @@ class TestNonAdminCannotPersistWifFieldsOnModel: model_id="m1", patch_data=updateDeployment( litellm_params=updateLiteLLMParams( - anthropic_keycloak_token_url="https://keycloak.internal/token", + **{field: value}, ) ), user_api_key_dict=admin, diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 37028558af3..f727448f3ac 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -54,6 +54,8 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _LOCK_DELETED_TEAM_MEMBERS_SQL, _STRIP_DELETED_TEAMS_FROM_USERS_SQL, DeletedTeams, + _enrich_team_list_with_caller_membership, + _TeamListCallerIdentity, GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, _build_team_list_where_conditions, @@ -91,6 +93,8 @@ from litellm.router import Router from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkTeamMemberAddRequest, BulkTeamMemberAddResponse, + TeamListCallerMembership, + TeamListItem, TeamMemberAddResult, ) from litellm.types.utils import StandardAuditLogPayload @@ -4172,6 +4176,7 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) with patch( @@ -4191,9 +4196,8 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): ) # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -4229,6 +4233,7 @@ async def test_list_team_v2_security_check_admin_user(): ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) # Should NOT raise an exception @@ -4242,9 +4247,8 @@ async def test_list_team_v2_security_check_admin_user(): ) # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -4310,10 +4314,8 @@ async def test_list_team_v2_with_status_deleted(): mock_db.litellm_teamtable.count.assert_not_called() # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 - assert len(result["teams"]) == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -4372,6 +4374,7 @@ async def test_list_team_v2_includes_litellm_model_table(): side_effect=lambda **kw: [_team_row("team_1", kw.get("include"))] ) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) result = await list_team_v2( @@ -4383,8 +4386,8 @@ async def test_list_team_v2_includes_litellm_model_table(): status=None, ) - assert result["teams"][0].litellm_model_table is not None - assert result["teams"][0].litellm_model_table.model_aliases == {"my-fast-model": "fake-model"} + assert result.teams[0].litellm_model_table is not None + assert result.teams[0].litellm_model_table.model_aliases == {"my-fast-model": "fake-model"} mock_db.litellm_deletedteamtable.find_many = AsyncMock( side_effect=lambda **kw: [_team_row("team_2", kw.get("include"))] @@ -4465,6 +4468,7 @@ async def test_list_team_v2_org_admin_sees_org_teams(): } mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) @@ -4482,9 +4486,9 @@ async def test_list_team_v2_org_admin_sees_org_teams(): status=None, ) - assert result["total"] == 1 - assert len(result["teams"]) == 1 - assert result["teams"][0].members_count == 1 + assert result.total == 1 + assert len(result.teams) == 1 + assert result.teams[0].members_count == 1 # Verify org-scoped where clause where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] @@ -4565,6 +4569,7 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): return_value=[mock_team_1, mock_team_2] ) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) @@ -4583,8 +4588,8 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): status=None, ) - assert result["total"] == 2 - assert len(result["teams"]) == 2 + assert result.total == 2 + assert len(result.teams) == 2 # Verify the where clause scopes by org OR own membership — no # top-level team_id filter that would hide org teams they aren't in @@ -4689,6 +4694,7 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( prisma_client = MagicMock() prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) prisma_client.db.litellm_teamtable.count = AsyncMock(side_effect=count) + prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) prisma_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=LiteLLM_UserTable( @@ -4721,8 +4727,8 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( sort_order="asc", status=None, ) - assert result["total"] == len(result["teams"]) - return [t.team_id for t in result["teams"]] + assert result.total == len(result.teams) + return [t.team_id for t in result.teams] own_view = ["team_in_org_A", "other_team_in_org_A", "team_in_org_B"] assert await list_teams("org_admin_user") == own_view @@ -4950,6 +4956,7 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): } mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) result = await list_team_v2( @@ -4966,7 +4973,7 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): status=None, ) - assert result["total"] == 1 + assert result.total == 1 # Verify the where clause filters by user's teams AND org scope where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] @@ -5219,6 +5226,7 @@ async def test_list_team_v2_populates_keys_count(): mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b]) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock( return_value=[ {"team_id": "team_a", "_count": {"team_id": 3}}, @@ -5235,16 +5243,391 @@ async def test_list_team_v2_populates_keys_count(): status=None, ) - assert result["total"] == 2 - by_id = {t.team_id: t for t in result["teams"]} + assert result.total == 2 + by_id = {t.team_id: t for t in result.teams} assert by_id["team_a"].keys_count == 3 assert by_id["team_b"].keys_count == 0 - # The aggregate is one batched query, filtered by the page's team IDs. group_by_kwargs = mock_db.litellm_verificationtoken.group_by.call_args.kwargs assert group_by_kwargs["by"] == ["team_id"] assert group_by_kwargs["where"] == {"team_id": {"in": ["team_a", "team_b"]}} assert group_by_kwargs["count"] == {"team_id": True} + membership_where: Final = mock_db.litellm_teammembership.find_many.call_args.kwargs["where"] + assert membership_where == { + "team_id": {"in": ("team_a", "team_b")}, + "user_id": "admin_user_123", + } + + +def _team_list_prisma( + membership_rows: Sequence[LiteLLM_TeamMembership], + budget_rows: Sequence[LiteLLM_BudgetTableFull], +) -> tuple[MagicMock, AsyncMock, AsyncMock]: + prisma_client: Final = MagicMock() + membership_query: Final = AsyncMock(return_value=tuple(membership_rows)) + budget_query: Final = AsyncMock(return_value=tuple(budget_rows)) + prisma_client.db.litellm_teammembership.find_many = membership_query + prisma_client.db.litellm_budgettable.find_many = budget_query + return prisma_client, membership_query, budget_query + + +def _team_list_member( + team_id: str, + user_id: str = "caller", + default_budget_id: str | None = None, +) -> TeamListItem: + metadata: Final = {"team_member_budget_id": default_budget_id} if default_budget_id is not None else None + return TeamListItem( + team_id=team_id, + metadata=metadata, + members_with_roles=[Member(user_id=user_id, role="user")], + ) + + +def _team_list_caller_identity( + user_id: str | None = "caller", + user_email: str | None = None, +) -> _TeamListCallerIdentity: + return _TeamListCallerIdentity(user_id=user_id, user_email=user_email) + + +def _budget_table( + budget_id: str, + max_budget: float | None, + budget_reset_at: datetime | None = None, + temp_budget_increase: float | None = None, + temp_budget_expiry: datetime | None = None, +) -> LiteLLM_BudgetTableFull: + return LiteLLM_BudgetTableFull( + budget_id=budget_id, + max_budget=max_budget, + budget_reset_at=budget_reset_at, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + temp_budget_increase=temp_budget_increase, + temp_budget_expiry=temp_budget_expiry, + ) + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_member_spend_budget_and_reset_date() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + member_budget: Final = _budget_table("member-budget", 100, reset_at) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + spend=50, + litellm_budget_table=member_budget, + ) + prisma_client, _, budget_query = _team_list_prisma([membership], []) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=50, + max_budget=100, + budget_reset_at=reset_at, + ) + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_includes_roster_members_without_membership_rows() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + prisma_client, membership_query, budget_query = _team_list_prisma( + (), + [_budget_table("team-budget", 50, reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=50, + budget_reset_at=None, + ) + membership_query.assert_awaited_once() + budget_query.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_matches_email_only_roster_member() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id=None, user_email="caller@x.com", role="user")], + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + (), + [_budget_table("team-budget", 50, reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + _team_list_caller_identity(user_email="caller@x.com"), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=50, + budget_reset_at=None, + ) + membership_query.assert_awaited_once() + budget_query.assert_awaited_once() + + +@pytest.mark.parametrize("caller_user_email", ("other@x.com", None)) +@pytest.mark.asyncio +async def test_team_list_caller_membership_ignores_email_only_roster_member_without_matching_email( + caller_user_email: str | None, +) -> None: + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id=None, user_email="caller@x.com", role="user")], + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + (), + [_budget_table("team-budget", 50)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + _team_list_caller_identity(user_email=caller_user_email), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_awaited_once() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_team_default_without_member_reset_date() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + membership: Final = LiteLLM_TeamMembership(user_id="caller", team_id="team-a", spend=12) + prisma_client, _, budget_query = _team_list_prisma( + [membership], + [_budget_table("team-budget", 50, reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=12, + max_budget=50, + budget_reset_at=None, + ) + budget_query.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_reset_date_from_member_linked_default_budget() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + member_budget: Final = _budget_table("team-budget", 50, reset_at) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + litellm_budget_table=member_budget, + ) + prisma_client, _, budget_query = _team_list_prisma([membership], []) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=50, + budget_reset_at=reset_at, + ) + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_ignores_zero_team_default_budget() -> None: + membership: Final = LiteLLM_TeamMembership(user_id="caller", team_id="team-a") + prisma_client, _, _ = _team_list_prisma( + [membership], + [_budget_table("team-budget", 0, datetime(2026, 10, 20, tzinfo=timezone.utc))], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership(spend=0.0, max_budget=None, budget_reset_at=None) + + +@pytest.mark.parametrize( + ("expiry", "expected_budget"), + ( + (datetime(2030, 1, 2, tzinfo=timezone.utc), 110), + (datetime(2029, 12, 31, tzinfo=timezone.utc), 100), + ), +) +@pytest.mark.asyncio +async def test_team_list_caller_membership_applies_only_active_temporary_increases( + expiry: datetime, + expected_budget: float, +) -> None: + member_budget: Final = _budget_table( + "member-budget", + 100, + datetime(2030, 2, 1, tzinfo=timezone.utc), + temp_budget_increase=10, + temp_budget_expiry=expiry, + ) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + litellm_budget_table=member_budget, + ) + prisma_client, _, _ = _team_list_prisma([membership], []) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + _team_list_caller_identity(), + datetime(2030, 1, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is not None + assert teams[0].caller_membership.max_budget == expected_budget + assert teams[0].caller_membership.budget_reset_at == datetime(2030, 2, 1, tzinfo=timezone.utc) + + +@pytest.mark.parametrize( + "member_reset_at", + ( + datetime(2030, 1, 15, tzinfo=timezone.utc), + None, + ), +) +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_private_row_reset_with_inherited_budget( + member_reset_at: datetime | None, +) -> None: + team_reset_at: Final = datetime(2030, 2, 1, tzinfo=timezone.utc) + member_budget: Final = _budget_table( + "member-budget", + None, + member_reset_at, + temp_budget_increase=10, + temp_budget_expiry=datetime(2030, 1, 2, tzinfo=timezone.utc), + ) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + litellm_budget_table=member_budget, + ) + prisma_client, _, _ = _team_list_prisma( + [membership], + [_budget_table("team-budget", 50, team_reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2030, 1, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=60, + budget_reset_at=member_reset_at, + ) + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_is_null_for_nonmember_admins() -> None: + prisma_client, membership_query, budget_query = _team_list_prisma([], []) + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id="another-user", role="user")], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_awaited_once() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_skips_queries_without_caller_user_id() -> None: + prisma_client, membership_query, budget_query = _team_list_prisma((), ()) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + _team_list_caller_identity(user_id=None), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_not_awaited() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_batches_page_membership_and_budget_queries() -> None: + memberships: Final = ( + LiteLLM_TeamMembership(user_id="caller", team_id="team-a"), + LiteLLM_TeamMembership(user_id="caller", team_id="team-b"), + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + memberships, + [_budget_table("budget-a", 50), _budget_table("budget-b", 75)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [ + _team_list_member("team-a", default_budget_id="budget-a"), + _team_list_member("team-b", default_budget_id="budget-b"), + ], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert [team.caller_membership for team in teams] == [ + TeamListCallerMembership(spend=0.0, max_budget=50, budget_reset_at=None), + TeamListCallerMembership(spend=0.0, max_budget=75, budget_reset_at=None), + ] + membership_query.assert_awaited_once() + budget_query.assert_awaited_once() @pytest.mark.asyncio @@ -5282,8 +5665,8 @@ async def test_list_team_v2_keys_count_skipped_for_empty_page(): status=None, ) - assert result["total"] == 0 - assert result["teams"] == [] + assert result.total == 0 + assert result.teams == [] mock_db.litellm_verificationtoken.group_by.assert_not_called() @@ -5332,7 +5715,7 @@ async def test_list_team_v2_keys_count_skipped_for_deleted_status(): status="deleted", ) - assert result["total"] == 1 + assert result.total == 1 mock_db.litellm_verificationtoken.group_by.assert_not_called() @@ -12178,6 +12561,36 @@ async def test_team_member_me_returns_caller_membership(mock_db_client): assert call_kwargs["team_id"] == team_id +@pytest.mark.asyncio +@pytest.mark.parametrize( + "spend, stored_total, reported_total", + [(612.4, 148.1, 612.4), (12.25, 97.8, 97.8)], + ids=["legacy_row_floors_to_cycle_spend", "healthy_row_keeps_its_total"], +) +async def test_team_member_me_reports_total_spend_no_lower_than_cycle_spend( + mock_db_client, spend, stored_total, reported_total +): + from fastapi import Request + + from litellm.proxy.management_endpoints.team_endpoints import team_member_me + + team_id = "team-me-floor" + caller_id = "alice@example.com" + team = _build_team_for_me(team_id, [{"user_id": caller_id, "user_email": None, "role": "user"}]) + membership = _build_membership_for_me(caller_id, team_id, spend=spend).model_copy( + update={"total_spend": stored_total} + ) + p_team, p_membership, p_user = _patch_member_me_helpers(team=team, membership=membership, user=None) + with p_team, p_membership, p_user: + response = await team_member_me( + http_request=MagicMock(spec=Request), + team_id=team_id, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id), + ) + + assert (response.spend, response.total_spend) == (spend, reported_total) + + @pytest.mark.asyncio async def test_team_member_me_matches_email_only_member(mock_db_client): """ @@ -16225,7 +16638,7 @@ async def test_reset_team_member_spend_fn_success(monkeypatch): litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0), ) mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) - mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + mock_prisma_client.db.execute_raw = AsyncMock(return_value=1) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) @@ -16248,15 +16661,76 @@ async def test_reset_team_member_spend_fn_success(monkeypatch): assert response["spend"] == 0.0 assert response["previous_spend"] == 10.0 assert response["max_budget"] == 50.0 - mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, - data={"spend": 0.0}, - ) + statement, reset_user_id, reset_team_id, reset_to = mock_prisma_client.db.execute_raw.await_args.args + assert (reset_user_id, reset_team_id, reset_to) == ("member-1", "team-1", 0.0) + assert "total_spend = GREATEST(total_spend, spend)" in statement + assert "spend >= $3" in statement assert await real_cache.async_get_cache(key="team-1_member-1") is None assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0 +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_stale_spend_returns_409_and_keeps_caches(monkeypatch): + """The reset validates reset_to against a spend it read a moment ago. If the row moved + underneath it (a flush landed, another admin reset first) the guarded UPDATE matches no + row, and the endpoint must say so instead of reporting a reset that never happened.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="cached-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0) + + membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.execute_raw = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ), pytest.raises(HTTPException) as exc_info: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="member-1", + data=ResetSpendRequest(reset_to=5.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert exc_info.value.status_code == 409 + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") == "cached-membership" + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0 + + +@pytest.mark.asyncio +async def test_get_all_team_memberships_floors_total_spend_at_current_cycle_spend(): + """Rows written before total_spend existed carry their history only in spend. /team/info + and /team/list both read through this helper, so a legacy row reports its lifetime spend + as at least the current cycle, while a healthy row keeps its own total.""" + from litellm.proxy.management_endpoints.team_endpoints import get_all_team_memberships + + legacy_row = MagicMock() + legacy_row.model_dump = lambda: {"user_id": "legacy", "team_id": "team-1", "spend": 612.4, "total_spend": 148.1} + healthy_row = MagicMock() + healthy_row.model_dump = lambda: {"user_id": "healthy", "team_id": "team-1", "spend": 12.25, "total_spend": 97.8} + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[legacy_row, healthy_row]) + + legacy, healthy = await get_all_team_memberships(mock_prisma_client, ["team-1"]) + + assert (legacy.spend, legacy.total_spend) == (612.4, 612.4) + assert (healthy.spend, healthy.total_spend) == (12.25, 97.8) + + @pytest.mark.asyncio async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch): mock_prisma_client = MagicMock() @@ -16372,7 +16846,7 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0) mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) - mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + mock_prisma_client.db.execute_raw = AsyncMock(return_value=1) with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests "litellm.proxy.management_endpoints.team_endpoints.get_team_object", @@ -17611,6 +18085,7 @@ async def test_update_team_team_admin_resending_stored_member_defaults_is_not_an disable_audit_logging_for_mocked_team, ): import contextlib + from litellm.models.budget import LiteLLM_BudgetTable with contextlib.ExitStack() as stack: diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 479391e54fe..e1a939273ae 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -9198,58 +9198,6 @@ async def test_saml_callback_blocked_when_admin_ui_disabled(): assert "Admin UI is Disabled" in response.body.decode() -@pytest.mark.asyncio -async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): - """An IdP-initiated assertion must not bypass the >5 free-SSO-user Enterprise gate - that /sso/key/generate enforces; the ACS re-checks it after validating the assertion, - so the entitlement DB query never runs on unvalidated input.""" - from litellm.proxy._types import ProxyException - from litellm.proxy.management_endpoints.types import CustomOpenID - from litellm.proxy.management_endpoints.ui_sso import saml_callback - - call_order: list[str] = [] - - async def _fake_handle_acs(**kwargs): - call_order.append("validate") - return (CustomOpenID( - id="dana@litellm.ai", - email="dana@litellm.ai", - first_name=None, - last_name=None, - display_name="dana", - picture=None, - provider="saml", - team_ids=[], - user_role=None, - ), None) - - async def _fake_count_billable_users(): - call_order.append("count") - return 6 - - async def _stream(): - yield b"SAMLResponse=signed-response" - - request_double = SimpleNamespace(cookies={}, headers={}, stream=_stream) - - with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}), patch( - "litellm.proxy.proxy_server.premium_user", False - ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( - "litellm.proxy.proxy_server.master_key", MASTER_KEY - ), patch( - "litellm.proxy.management_endpoints.sso.saml_sso.SAMLAuthHandler.handle_acs", - new=_fake_handle_acs, - ), patch( - "litellm.repositories.user_repository.UserRepository.count_billable_users", - new=AsyncMock(side_effect=_fake_count_billable_users), - ): - with pytest.raises(ProxyException) as exc: - await saml_callback(request_double) - - assert str(exc.value.code) == "403" - assert call_order == ["validate", "count"] - - @pytest.mark.asyncio async def test_cli_poll_key_tolerates_missing_user_row(): """The CLI poll must still mint the JWT when the user lookup raises, diff --git a/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py deleted file mode 100644 index 4daddb89350..00000000000 --- a/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ /dev/null @@ -1,657 +0,0 @@ -""" -Tests for AI Usage Chat module. -""" - -import json -from collections.abc import AsyncIterator -from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from fastapi import HTTPException - -from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( - TOOL_HANDLERS, - TOOLS_ADMIN, - TOOLS_BASE, - _build_system_prompt, - _summarise_entity_data, - _summarise_usage_data, - stream_usage_ai_chat, -) - - -SAMPLE_AGGREGATED_RESPONSE = { - "results": [ - { - "date": "2025-01-15", - "metrics": { - "spend": 50.25, - "prompt_tokens": 20000, - "completion_tokens": 10000, - "total_tokens": 30000, - "api_requests": 500, - "successful_requests": 480, - "failed_requests": 20, - "cache_read_input_tokens": 0, - "cache_creation_input_tokens": 0, - }, - "breakdown": { - "models": { - "gpt-4": { - "metrics": { - "spend": 40.0, - "api_requests": 300, - "total_tokens": 25000, - }, - "metadata": {}, - "api_key_breakdown": {}, - }, - }, - "providers": { - "openai": { - "metrics": {"spend": 50.25, "api_requests": 500}, - "metadata": {}, - "api_key_breakdown": {}, - }, - }, - "api_keys": { - "sk-test123": { - "metrics": {"spend": 50.25}, - "metadata": {"key_alias": "Production Key"}, - }, - }, - "model_groups": {}, - "mcp_servers": {}, - "entities": {}, - }, - }, - ], - "metadata": { - "total_spend": 50.25, - "total_api_requests": 500, - "total_successful_requests": 480, - "total_failed_requests": 20, - "total_tokens": 30000, - }, -} - -SAMPLE_TEAM_RESPONSE = { - "results": [ - { - "date": "2025-01-15", - "metrics": {"spend": 100.0, "api_requests": 1000, "total_tokens": 50000}, - "breakdown": { - "entities": { - "team-1": { - "metrics": { - "spend": 60.0, - "api_requests": 600, - "total_tokens": 30000, - }, - "metadata": {"alias": "Engineering"}, - "api_key_breakdown": {}, - }, - "team-2": { - "metrics": { - "spend": 40.0, - "api_requests": 400, - "total_tokens": 20000, - }, - "metadata": {"alias": "Marketing"}, - "api_key_breakdown": {}, - }, - }, - "models": {}, - "providers": {}, - "api_keys": {}, - "model_groups": {}, - "mcp_servers": {}, - }, - }, - ], - "metadata": {"total_spend": 100.0, "total_api_requests": 1000}, -} - - -class TestToolSchemas: - def test_admin_tools_include_all(self): - assert len(TOOLS_ADMIN) == 3 - names = {t["function"]["name"] for t in TOOLS_ADMIN} - assert "get_usage_data" in names - assert "get_team_usage_data" in names - assert "get_tag_usage_data" in names - - def test_base_tools_restricted_to_usage_only(self): - assert len(TOOLS_BASE) == 1 - assert TOOLS_BASE[0]["function"]["name"] == "get_usage_data" - - def test_admin_prompt_mentions_all_tools(self): - prompt = _build_system_prompt(is_admin=True) - assert "get_usage_data" in prompt - assert "get_team_usage_data" in prompt - assert "get_tag_usage_data" in prompt - - def test_non_admin_prompt_only_mentions_usage_tool(self): - prompt = _build_system_prompt(is_admin=False) - assert "get_usage_data" in prompt - assert "get_team_usage_data" not in prompt - assert "get_tag_usage_data" not in prompt - - def test_system_prompt_includes_todays_date(self): - from datetime import date - - prompt = _build_system_prompt(is_admin=True) - assert date.today().isoformat() in prompt - - -class TestSummariseUsageData: - def test_summarise_includes_totals(self): - summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) - assert "$50.25" in summary - assert "500" in summary - - def test_summarise_includes_models(self): - summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) - assert "gpt-4" in summary - - def test_summarise_includes_providers(self): - summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) - assert "openai" in summary - - def test_summarise_handles_empty_data(self): - empty = {"results": [], "metadata": {}} - summary = _summarise_usage_data(empty) - assert "no data" in summary.lower() - - -class TestSummariseEntityData: - def test_team_summary_includes_teams(self): - summary = _summarise_entity_data(SAMPLE_TEAM_RESPONSE, "Team") - assert "Engineering" in summary - assert "Marketing" in summary - assert "$60.0" in summary - assert "$40.0" in summary - - def test_team_summary_empty(self): - empty = {"results": [], "metadata": {}} - summary = _summarise_entity_data(empty, "Team") - assert "No Team usage data" in summary - - -class TestStreamUsageAiChat: - @pytest.mark.asyncio - async def test_stream_emits_status_events(self): - mock_tool_call = MagicMock() - mock_tool_call.id = "call_123" - mock_tool_call.function.name = "get_usage_data" - mock_tool_call.function.arguments = json.dumps( - { - "start_date": "2025-01-01", - "end_date": "2025-01-31", - } - ) - - mock_first_response = MagicMock() - mock_first_response.choices = [MagicMock()] - mock_first_response.choices[0].message.tool_calls = [mock_tool_call] - mock_first_response.choices[0].message.model_dump.return_value = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_123", - "type": "function", - "function": { - "name": "get_usage_data", - "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}', - }, - } - ], - } - - async def mock_stream(): - chunk = MagicMock() - chunk.choices = [MagicMock()] - chunk.choices[0].delta.content = "Total spend is $50.25" - yield chunk - - with ( - patch( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" - ) as mock_litellm, - patch( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_usage_data", - new_callable=AsyncMock, - ) as mock_fetch, - ): - mock_litellm.acompletion = AsyncMock( - side_effect=[ - mock_first_response, - mock_stream(), - ] - ) - mock_fetch.return_value = SAMPLE_AGGREGATED_RESPONSE - - events = [] - async for event in stream_usage_ai_chat( - messages=[{"role": "user", "content": "What is my total spend?"}], - model="gpt-4o-mini", - user_id="user-123", - is_admin=True, - ): - events.append(json.loads(event.replace("data: ", "").strip())) - - status_events = [e for e in events if e["type"] == "status"] - tool_call_events = [e for e in events if e["type"] == "tool_call"] - chunk_events = [e for e in events if e["type"] == "chunk"] - done_events = [e for e in events if e["type"] == "done"] - - assert len(status_events) >= 1 - assert "Thinking" in status_events[0]["message"] - assert len(tool_call_events) >= 1 - assert tool_call_events[0]["tool_name"] == "get_usage_data" - assert tool_call_events[0]["status"] in ("running", "complete") - assert len(chunk_events) >= 1 - assert len(done_events) == 1 - - @pytest.mark.asyncio - async def test_stream_handles_team_tool(self): - mock_tool_call = MagicMock() - mock_tool_call.id = "call_team" - mock_tool_call.function.name = "get_team_usage_data" - mock_tool_call.function.arguments = json.dumps( - { - "start_date": "2025-01-01", - "end_date": "2025-01-31", - } - ) - - mock_first_response = MagicMock() - mock_first_response.choices = [MagicMock()] - mock_first_response.choices[0].message.tool_calls = [mock_tool_call] - mock_first_response.choices[0].message.model_dump.return_value = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_team", - "type": "function", - "function": { - "name": "get_team_usage_data", - "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}', - }, - } - ], - } - - async def mock_stream(): - chunk = MagicMock() - chunk.choices = [MagicMock()] - chunk.choices[0].delta.content = "Engineering is the top team." - yield chunk - - with ( - patch( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" - ) as mock_litellm, - patch( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_team_usage_data", - new_callable=AsyncMock, - ) as mock_fetch, - ): - mock_litellm.acompletion = AsyncMock( - side_effect=[ - mock_first_response, - mock_stream(), - ] - ) - mock_fetch.return_value = SAMPLE_TEAM_RESPONSE - - events = [] - async for event in stream_usage_ai_chat( - messages=[{"role": "user", "content": "Which team spends the most?"}], - model="gpt-4o-mini", - is_admin=True, - ): - events.append(json.loads(event.replace("data: ", "").strip())) - - chunk_events = [e for e in events if e["type"] == "chunk"] - assert len(chunk_events) >= 1 - assert "Engineering" in chunk_events[0]["content"] - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "reason", - [ - "end_date must be on or after start_date", - "start_date and end_date must be valid YYYY-MM-DD dates", - ], - ids=["reversed", "invalid-format"], - ) - @pytest.mark.parametrize( - "tool_name", ["get_usage_data", "get_team_usage_data", "get_tag_usage_data"] - ) - async def test_invalid_tool_dates_are_reported_to_the_model( - self, tool_name: str, reason: str - ) -> None: - fetch: Final = AsyncMock( - side_effect=HTTPException(status_code=400, detail={"error": reason}) - ) - - tool_message, statuses = await _run_one_tool( - tool_name, "2025-01-31", "2025-01-01", fetch - ) - - fetch.assert_awaited_once() - assert tool_message == f"Invalid date range: {reason}" - assert statuses == ["running", "complete"] - - @pytest.mark.asyncio - async def test_tool_dates_the_query_layer_accepts_are_queried_as_sent( - self, - ) -> None: - fetch: Final = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE) - - tool_message, statuses = await _run_one_tool( - "get_usage_data", "2025/01/01", "2025-01-31T00:00:00", fetch - ) - - fetch.assert_awaited_once_with( - start_date="2025/01/01", end_date="2025-01-31T00:00:00" - ) - assert tool_message.startswith("Total Spend: $") - assert statuses == ["running", "complete"] - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "error", - [ - HTTPException(status_code=500, detail={"error": "db down"}), - HTTPException(status_code=400, detail="page and page_size must be >= 1"), - ], - ids=["server-error", "other-400"], - ) - async def test_other_tool_errors_stay_a_generic_fetch_error( - self, error: HTTPException - ) -> None: - tool_message, statuses = await _run_one_tool( - "get_team_usage_data", - "2025-01-01", - "2025-01-31", - AsyncMock(side_effect=error), - ) - - assert tool_message == "Error fetching team usage data. Please try again." - assert statuses == ["running", "error"] - - @pytest.mark.asyncio - async def test_stream_handles_error(self): - with patch( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" - ) as mock_litellm: - mock_litellm.acompletion = AsyncMock(side_effect=Exception("LLM error")) - - events = [] - async for event in stream_usage_ai_chat( - messages=[{"role": "user", "content": "test"}], - ): - events.append(json.loads(event.replace("data: ", "").strip())) - - error_events = [e for e in events if e["type"] == "error"] - assert len(error_events) == 1 - assert "internal error" in error_events[0]["message"].lower() - - @pytest.mark.asyncio - async def test_non_admin_enforces_user_id(self): - mock_tool_call = MagicMock() - mock_tool_call.id = "call_456" - mock_tool_call.function.name = "get_usage_data" - mock_tool_call.function.arguments = json.dumps( - { - "start_date": "2025-01-01", - "end_date": "2025-01-31", - "user_id": "other-user", - } - ) - - mock_first_response = MagicMock() - mock_first_response.choices = [MagicMock()] - mock_first_response.choices[0].message.tool_calls = [mock_tool_call] - mock_first_response.choices[0].message.model_dump.return_value = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_456", - "type": "function", - "function": { - "name": "get_usage_data", - "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31","user_id":"other-user"}', - }, - } - ], - } - - async def mock_stream(): - chunk = MagicMock() - chunk.choices = [MagicMock()] - chunk.choices[0].delta.content = "Data." - yield chunk - - mock_fetch = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE) - - with ( - patch( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" - ) as mock_litellm, - patch.dict( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.TOOL_HANDLERS", - { - "get_usage_data": { - "fetch": mock_fetch, - "summarise": _summarise_usage_data, - "label": "global usage data", - } - }, - ), - ): - mock_litellm.acompletion = AsyncMock( - side_effect=[ - mock_first_response, - mock_stream(), - ] - ) - - events = [] - async for event in stream_usage_ai_chat( - messages=[{"role": "user", "content": "Show data"}], - model="gpt-4o-mini", - user_id="my-user-id", - is_admin=False, - ): - events.append(event) - - mock_fetch.assert_called_once_with( - start_date="2025-01-01", - end_date="2025-01-31", - user_id="my-user-id", - ) - - -class TestUsageAiChatServiceAccountGuard: - """ - Security regression: a non-admin caller with user_id=None (service-account - key) must be rejected at the endpoint boundary, before any tool dispatch. - """ - - @pytest.mark.asyncio - async def test_non_admin_with_user_id_none_is_rejected(self): - from fastapi import HTTPException - - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( - ChatMessage, - UsageAIChatRequest, - usage_ai_chat, - ) - - service_account_key = UserAPIKeyAuth( - user_id=None, - user_role=LitellmUserRoles.INTERNAL_USER, - ) - request = MagicMock() - body = UsageAIChatRequest( - messages=[ChatMessage(role="user", content="hi")], - model="gpt-4o-mini", - ) - - with pytest.raises(HTTPException) as exc_info: - await usage_ai_chat( - data=body, - request=request, - user_api_key_dict=service_account_key, - ) - - assert exc_info.value.status_code == 403 - assert "Service-account keys" in str(exc_info.value.detail) - - def test_resolve_fetch_kwargs_tripwire_fires_on_none_user_id(self): - """ - Defense-in-depth: if a future endpoint forgets the entry guard and - a non-admin caller with user_id=None reaches _resolve_fetch_kwargs, - the tripwire must fire rather than issuing an unscoped query. - """ - from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( - _resolve_fetch_kwargs, - ) - - with pytest.raises(ValueError, match='Non-admin caller has user_id=None; refusing to issue an') as exc_info: - _resolve_fetch_kwargs( - fn_name="get_usage_data", - fn_args={"start_date": "2025-01-01", "end_date": "2025-01-31"}, - user_id=None, - is_admin=False, - ) - assert "Endpoint-level guard missing" in str(exc_info.value) - - -class TestUsageAiChatKeepalive: - async def _collect_endpoint_body(self, monkeypatch, interval, delay=0.3) -> tuple[list[bytes], dict]: - import asyncio - - import litellm - from fastapi.responses import StreamingResponse - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( - ChatMessage, - UsageAIChatRequest, - usage_ai_chat, - ) - - monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) - - async def slow_acompletion(**kwargs): - await asyncio.sleep(delay) - response = MagicMock() - response.choices = [MagicMock()] - response.choices[0].message.tool_calls = None - response.choices[0].message.content = "Total spend is $50.25" - return response - - with patch( # test-quality-ok: the stream calls the module-level litellm.acompletion directly; no injection seam - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm.acompletion", - new=AsyncMock(side_effect=slow_acompletion), - ): - response = await usage_ai_chat( - data=UsageAIChatRequest(messages=[ChatMessage(role="user", content="hi")], model="gpt-4o-mini"), - request=MagicMock(), - user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), - ) - assert isinstance(response, StreamingResponse) - chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] - return chunks, dict(response.headers) - - @pytest.mark.asyncio - async def test_endpoint_pings_while_the_planning_completion_is_still_running(self, monkeypatch): - chunks, headers = await self._collect_endpoint_body(monkeypatch, interval=0.05) - - assert headers["content-type"].startswith("text/event-stream") - assert headers["cache-control"] == "no-cache" - assert headers["x-accel-buffering"] == "no" - assert chunks[0].startswith(b'data: {"type": "status"') - assert chunks[1] == b": ping\n\n" - assert chunks.count(b": ping\n\n") >= 3 - assert b'"content": "Total spend is $50.25"' in b"".join(chunks) - assert chunks[-1] == b'data: {"type": "done"}\n\n' - - @pytest.mark.asyncio - async def test_endpoint_stream_is_untouched_while_keepalives_are_unconfigured(self, monkeypatch): - chunks, _ = await self._collect_endpoint_body(monkeypatch, interval=None, delay=0.15) - - assert b": ping\n\n" not in chunks - assert chunks[0].startswith(b'data: {"type": "status"') - assert chunks[-1] == b'data: {"type": "done"}\n\n' - - -async def _run_one_tool( - tool_name: str, start_date: str, end_date: str, fetch: AsyncMock -) -> tuple[str, list[object]]: - """Runs one admin chat turn whose model calls `tool_name` with these dates; returns the - tool message the follow-up model call carried and the tool statuses streamed.""" - arguments: Final = json.dumps({"start_date": start_date, "end_date": end_date}) - mock_tool_call: Final = MagicMock() - mock_tool_call.id = "call_dates" - mock_tool_call.function.name = tool_name - mock_tool_call.function.arguments = arguments - - mock_first_response: Final = MagicMock() - mock_first_response.choices = [MagicMock()] - mock_first_response.choices[0].message.tool_calls = [mock_tool_call] - mock_first_response.choices[0].message.model_dump.return_value = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_dates", - "type": "function", - "function": {"name": tool_name, "arguments": arguments}, - } - ], - } - - async def mock_stream() -> AsyncIterator[MagicMock]: - chunk: Final = MagicMock() - chunk.choices = [MagicMock()] - chunk.choices[0].delta.content = "Done." - yield chunk - - with ( - patch( - "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" - ) as mock_litellm, - patch.dict(TOOL_HANDLERS[tool_name], {"fetch": fetch}), - ): - mock_litellm.acompletion = AsyncMock( - side_effect=[mock_first_response, mock_stream()] - ) - events: Final[tuple[dict[str, object], ...]] = tuple( - [ - json.loads(event.replace("data: ", "").strip()) - async for event in stream_usage_ai_chat( - messages=[{"role": "user", "content": "How much did we spend?"}], - model="gpt-4o-mini", - is_admin=True, - ) - ] - ) - final_messages: Final = mock_litellm.acompletion.call_args_list[1].kwargs[ - "messages" - ] - - tool_messages: Final = [m for m in final_messages if m["role"] == "tool"] - assert len(tool_messages) == 1, tool_messages - assert tool_messages[0]["tool_call_id"] == "call_dates" - return str(tool_messages[0]["content"]), [ - e["status"] for e in events if e["type"] == "tool_call" - ] diff --git a/tests/unit/proxy/management_helpers/test_team_list_sorting.py b/tests/unit/proxy/management_helpers/test_team_list_sorting.py new file mode 100644 index 00000000000..bc4c77cc6a4 --- /dev/null +++ b/tests/unit/proxy/management_helpers/test_team_list_sorting.py @@ -0,0 +1,85 @@ +import json +import re +import sqlite3 +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import JsonValue + +from litellm.proxy.management_helpers.team_list_sorting import UnsupportedTeamFilter, _where_sql, sorted_team_ids + + +@pytest.mark.parametrize( + ("where", "expected"), + ( + ({}, ("a", "b", "c")), + ({"OR": []}, ()), + ({"AND": []}, ("a", "b", "c")), + ({"team_id": "a"}, ("a",)), + ({"team_alias": {"contains": "ALPHA", "mode": "insensitive"}}, ("a", "c")), + ({"team_id": {"startsWith": "b"}}, ("b",)), + ({"organization_id": {"in": ["org-1"]}}, ("a", "b")), + ( + { + "OR": [{"team_id": "b"}, {"team_alias": {"contains": "alpha", "mode": "insensitive"}}], + "AND": [{"OR": [{"organization_id": "org-1"}, {"team_id": "c"}]}], + "team_id": {"in": ["a", "c"]}, + }, + ("a", "c"), + ), + ({"team_id": "a' OR 1=1 --"}, ()), + ), +) +def test_team_filters_preserve_search_and_authorization_intersection( + where: dict[str, JsonValue], expected: tuple[str, ...] +) -> None: + result: Final = _where_sql(where) + assert not isinstance(result, UnsupportedTeamFilter) + predicate, params = result + sqlite_predicate: Final = re.sub( + r"= ANY\((\$\d+)::text\[\]\)", r"IN (SELECT value FROM json_each(\1))", predicate + ).replace(" ILIKE ", " LIKE ") + with sqlite3.connect(":memory:") as db: + db.execute("CREATE TABLE teams (team_id TEXT, team_alias TEXT, organization_id TEXT)") + db.executemany( + "INSERT INTO teams VALUES (?, ?, ?)", + (("a", "Alpha", "org-1"), ("b", "Beta", "org-1"), ("c", "Alpha", "org-2")), + ) + rows: Final = db.execute( + f"SELECT t.team_id FROM teams t WHERE {sqlite_predicate} ORDER BY t.team_id", + { + str(index): json.dumps(value) if isinstance(value, list) else value + for index, value in enumerate(params, 1) + }, + ).fetchall() + assert tuple(row[0] for row in rows) == expected + + +@pytest.mark.parametrize("where", ({"unknown": "x"}, {"team_id": {"not": "x"}}, {"AND": [False]})) +@pytest.mark.asyncio +async def test_unsupported_filter_fails_without_reading_unscoped_teams(where: dict[str, JsonValue]) -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock() + result: Final = await sorted_team_ids(prisma, where=where, sort_by="keys_count", descending=False, skip=0, take=25) + assert isinstance(result, UnsupportedTeamFilter) + prisma.db.query_raw.assert_not_awaited() + + +@pytest.mark.parametrize("sort_by", ("keys_count", "members", "models", "organization_alias")) +@pytest.mark.parametrize("descending", (False, True)) +@pytest.mark.asyncio +async def test_sorted_ids_bind_search_and_pagination_without_interpolating_input( + sort_by: str, descending: bool +) -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + search: Final = "x' OR TRUE --" + result: Final = await sorted_team_ids( + prisma, where={"team_id": search}, sort_by=sort_by, descending=descending, skip=25, take=25 + ) + assert result == () + sql, *bindings = prisma.db.query_raw.await_args.args + assert search not in sql + assert bindings == [search, 25, 25] + assert ("DESC" if descending else "ASC") + " NULLS LAST, t.team_id ASC" in sql diff --git a/tests/unit/proxy/management_helpers/test_user_list_sorting.py b/tests/unit/proxy/management_helpers/test_user_list_sorting.py new file mode 100644 index 00000000000..263c1691b23 --- /dev/null +++ b/tests/unit/proxy/management_helpers/test_user_list_sorting.py @@ -0,0 +1,71 @@ +from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy.management_helpers.user_list_sorting import sorted_user_ids + + +@pytest.mark.parametrize("sort_by", ("status", "key_count")) +@pytest.mark.parametrize("descending", (False, True)) +@pytest.mark.asyncio +async def test_user_sort_binds_scope_search_and_page_and_excludes_ui_sessions( + sort_by: Literal["status", "key_count"], descending: bool +) -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + search: Final = "x' OR TRUE --" + result: Final = await sorted_user_ids( + prisma, + sort_by=sort_by, + descending=descending, + skip=25, + take=25, + role="internal_user", + user_ids=" a, b ", + user_email=search, + team="team-1", + sso_user_ids="sso-1, sso-2", + organization_ids="org-1, org-2", + search=search, + ) + assert result == () + sql, *bindings = prisma.db.query_raw.await_args.args + assert search not in sql + assert bindings == [ + "internal_user", + ("a", "b"), + search, + "team-1", + False, + ("sso-1", "sso-2"), + ("org-1", "org-2"), + search, + 25, + 25, + ] + ([UI_SESSION_TOKEN_TEAM_ID] if sort_by == "key_count" else []) + assert ("DESC" if descending else "ASC") + ", u.user_id ASC" in sql + assert ("k.team_id IS DISTINCT FROM $11::text" in sql) == (sort_by == "key_count") + + +@pytest.mark.asyncio +async def test_missing_user_filters_preserve_unfiltered_query() -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + result: Final = await sorted_user_ids( + prisma, + sort_by="status", + descending=False, + skip=0, + take=25, + role=None, + user_ids=None, + user_email=None, + team=None, + sso_user_ids=None, + organization_ids=None, + search=None, + ) + assert result == () + assert prisma.db.query_raw.await_args.args[1:] == (None, None, None, None, True, (), (), None, 25, 0) diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index eb222366812..d47db69257b 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -722,6 +722,16 @@ def _router_with_executed_batch_model() -> Router: "litellm_params": {"model": "gemini/gemini-2.0-flash"}, "model_info": {"id": "gemini-2.0-flash-id"}, }, + { + "model_name": "vertex-batch-drop", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "drop_params": True}, + "model_info": {"id": "vertex-batch-drop-id"}, + }, + { + "model_name": "vertex-batch-keep", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "drop_params": False}, + "model_info": {"id": "vertex-batch-keep-id"}, + }, ] ) @@ -939,6 +949,39 @@ def test_batch_upload_for_a_provider_model_still_goes_to_the_provider(batch_uplo assert provider_upload.call_args.kwargs["custom_llm_provider"] == "gemini" +@pytest.mark.parametrize( + "model, expected_drop_params", + [("vertex-batch-drop", True), ("vertex-batch-keep", False)], +) +def test_batch_upload_for_a_provider_model_carries_the_deployment_drop_params( + batch_upload_seams, model: str, expected_drop_params: bool +): + _, provider_upload, _ = batch_upload_seams + + response = _upload_batch_file({"x-litellm-model": model}, {}) + + assert response.status_code == 200, response.text + assert provider_upload.call_args.kwargs["drop_params"] is expected_drop_params + + +def test_batch_upload_for_a_provider_model_without_deployment_drop_params_sends_no_drop_params(batch_upload_seams): + _, provider_upload, _ = batch_upload_seams + + response = _upload_batch_file({"x-litellm-model": "gemini-2.0-flash"}, {}) + + assert response.status_code == 200, response.text + assert "drop_params" not in provider_upload.call_args.kwargs + + +def test_batch_upload_keeps_the_callers_drop_params_over_the_deployment(batch_upload_seams): + _, provider_upload, _ = batch_upload_seams + + response = _upload_batch_file({"x-litellm-model": "vertex-batch-keep", "user-agent": "claude-cli/2.1.0"}, {}) + + assert response.status_code == 200, response.text + assert provider_upload.call_args.kwargs["drop_params"] is True + + @pytest.mark.skip(reason="mock respx fails on ci/cd - unclear why") def test_create_file_and_call_chat_completion_e2e( mocker: MockerFixture, monkeypatch, llm_router: Router diff --git a/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py b/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py index 81c5803da33..09104c035c7 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py +++ b/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import MagicMock import pytest @@ -27,10 +28,7 @@ class _RecordingStorageBackend: raise self.delete_error -class _FakeManagedFilesHook(BaseFileEndpoints): - def __init__(self): - self.stored = [] - +class _ManagedFilesEndpoints(BaseFileEndpoints): async def acreate_file( self, create_file_request, llm_router, target_model_names_list, litellm_parent_otel_span, user_api_key_dict ): @@ -48,10 +46,23 @@ class _FakeManagedFilesHook(BaseFileEndpoints): async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): raise NotImplementedError + +class _FakeManagedFilesHook(_ManagedFilesEndpoints): + def __init__(self): + self.stored = [] + async def store_unified_file_id(self, **kwargs): self.stored.append(kwargs) +class _ForwardingManagedFilesHook(_ManagedFilesEndpoints): + def __init__(self, writer: _FakeManagedFilesHook): + self._writer = writer + + def __getattr__(self, name: str) -> object: + return getattr(self._writer, name) + + class _FailingManagedFilesHook(_FakeManagedFilesHook): async def store_unified_file_id(self, **kwargs): raise RuntimeError("db down") @@ -142,6 +153,24 @@ async def test_upload_with_target_model_names_and_hook_stores_unified_id(monkeyp } +@pytest.mark.asyncio +async def test_upload_stores_the_unified_id_through_a_hook_that_forwards_its_writer(monkeypatch): + backend: Final = _RecordingStorageBackend() + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name, prisma_client=None: backend) + writer: Final = _FakeManagedFilesHook() + + file_object: Final = await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=["gpt-x"], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=_ForwardingManagedFilesHook(writer)), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert [stored["file_id"] for stored in writer.stored] == [file_object.id] + + @pytest.mark.asyncio async def test_upload_hands_the_prisma_client_to_the_storage_backend_factory(monkeypatch: pytest.MonkeyPatch): backend = _RecordingStorageBackend() diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 4770ecd5485..7530703a05a 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -145,9 +145,7 @@ class TestAnthropicLoggingHandlerModelFallback: '{"type": "message_stop"}', ] - def _create_mock_logging_obj( - self, model_in_details: str = None - ) -> LiteLLMLoggingObj: + def _create_mock_logging_obj(self, model_in_details: str = None) -> LiteLLMLoggingObj: """Create a mock logging object with optional model in model_call_details""" mock_logging_obj = MagicMock() @@ -166,21 +164,13 @@ class TestAnthropicLoggingHandlerModelFallback: mock_handler = MagicMock() return mock_handler - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) - @patch.object( - AnthropicPassthroughLoggingHandler, "_create_anthropic_response_logging_payload" - ) - def test_model_from_request_body_used_when_present( - self, mock_create_payload, mock_build_response - ): + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") + @patch.object(AnthropicPassthroughLoggingHandler, "_create_anthropic_response_logging_payload") + def test_model_from_request_body_used_when_present(self, mock_create_payload, mock_build_response): """Test that model from request_body is used when present""" # Arrange request_body = {"model": "claude-3-sonnet-20240229"} - logging_obj = self._create_mock_logging_obj( - model_in_details="claude-3-haiku-20240307" - ) + logging_obj = self._create_mock_logging_obj(model_in_details="claude-3-haiku-20240307") passthrough_handler = self._create_mock_passthrough_handler() # Mock successful response building @@ -204,57 +194,37 @@ class TestAnthropicLoggingHandlerModelFallback: # Verify that _build_complete_streaming_response was called with the request_body model mock_build_response.assert_called_once() call_args = mock_build_response.call_args - assert ( - call_args[1]["model"] == "claude-3-sonnet-20240229" - ) # Should use request_body model + assert call_args[1]["model"] == "claude-3-sonnet-20240229" # Should use request_body model def test_model_fallback_logic_isolated(self): """Test just the model fallback logic in isolation""" # Test case 1: Model from request body request_body = {"model": "claude-3-sonnet-20240229"} - logging_obj = self._create_mock_logging_obj( - model_in_details="claude-3-haiku-20240307" - ) + logging_obj = self._create_mock_logging_obj(model_in_details="claude-3-haiku-20240307") # Extract the logic directly from the function model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "claude-3-sonnet-20240229" # Should use request_body model # Test case 2: Fallback to logging obj request_body = {} - logging_obj = self._create_mock_logging_obj( - model_in_details="claude-3-haiku-20240307" - ) + logging_obj = self._create_mock_logging_obj(model_in_details="claude-3-haiku-20240307") model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "claude-3-haiku-20240307" # Should use fallback model # Test case 3: Empty string in request body, fallback to logging obj request_body = {"model": ""} - logging_obj = self._create_mock_logging_obj( - model_in_details="claude-3-opus-20240229" - ) + logging_obj = self._create_mock_logging_obj(model_in_details="claude-3-opus-20240229") model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "claude-3-opus-20240229" # Should use fallback model @@ -264,11 +234,7 @@ class TestAnthropicLoggingHandlerModelFallback: logging_obj = self._create_mock_logging_obj() model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "" # Should be empty @@ -284,11 +250,7 @@ class TestAnthropicLoggingHandlerModelFallback: # Extract the logic directly from the function model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "" # Should remain empty since no fallback available @@ -298,11 +260,7 @@ class TestAnthropicLoggingHandlerModelFallback: logging_obj = self._create_mock_logging_obj() # Empty dict model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "" # Should remain empty @@ -311,9 +269,7 @@ class TestAnthropicLoggingHandlerModelFallback: class TestAzureAnthropicCostCalculation: """Test the custom_llm_provider cost calculation logic for Azure AI Anthropic.""" - def _create_mock_logging_obj( - self, model: str = None, custom_llm_provider: str = None - ) -> LiteLLMLoggingObj: + def _create_mock_logging_obj(self, model: str = None, custom_llm_provider: str = None) -> LiteLLMLoggingObj: """Create a mock logging object with optional model and custom_llm_provider""" mock_logging_obj = MagicMock() mock_model_call_details = {} @@ -326,9 +282,7 @@ class TestAzureAnthropicCostCalculation: return mock_logging_obj @patch("litellm.completion_cost") - def test_cost_calculation_with_azure_ai_custom_llm_provider( - self, mock_completion_cost - ): + def test_cost_calculation_with_azure_ai_custom_llm_provider(self, mock_completion_cost): """Test that custom_llm_provider is passed to completion_cost for Azure AI Anthropic""" from litellm.types.utils import ModelResponse from datetime import datetime @@ -397,9 +351,7 @@ class TestAzureAnthropicCostCalculation: assert call_kwargs["custom_llm_provider"] is None @patch("litellm.completion_cost") - def test_cost_calculation_does_not_duplicate_provider_prefix( - self, mock_completion_cost - ): + def test_cost_calculation_does_not_duplicate_provider_prefix(self, mock_completion_cost): """Test that provider prefix is not duplicated if already present in model name""" from litellm.types.utils import ModelResponse from datetime import datetime @@ -436,9 +388,7 @@ class TestAzureAnthropicCostCalculation: assert call_kwargs["custom_llm_provider"] == "azure_ai" @patch("litellm.completion_cost") - def test_cost_calculation_resolves_unknown_model_from_litellm_params( - self, mock_completion_cost - ): + def test_cost_calculation_resolves_unknown_model_from_litellm_params(self, mock_completion_cost): """When the body model is the "unknown" sentinel, the deployment model from litellm_params must be used for costing, not "unknown" (which makes completion_cost raise and the cost silently fall back to $0).""" @@ -451,9 +401,7 @@ class TestAzureAnthropicCostCalculation: logging_obj = self._create_mock_logging_obj(model="unknown") logging_obj.model_call_details["litellm_params"] = { "model": "anthropic/claude-3-5-haiku-20241022", - "metadata": { - "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" - }, + "metadata": {"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"}, } logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] @@ -471,17 +419,12 @@ class TestAzureAnthropicCostCalculation: ) mock_completion_cost.assert_called_once() - assert ( - mock_completion_cost.call_args[1]["model"] - == "anthropic/claude-3-5-haiku-20241022" - ) + assert mock_completion_cost.call_args[1]["model"] == "anthropic/claude-3-5-haiku-20241022" assert kwargs["response_cost"] == 0.001 assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022" @patch("litellm.completion_cost") - def test_cost_calculation_resolves_unknown_model_from_model_group( - self, mock_completion_cost - ): + def test_cost_calculation_resolves_unknown_model_from_model_group(self, mock_completion_cost): """With only model_group available (no deployment litellm_params.model), the leading passthrough/ prefix must be stripped so the cost map can resolve the model.""" @@ -493,9 +436,7 @@ class TestAzureAnthropicCostCalculation: logging_obj = self._create_mock_logging_obj(model="unknown") logging_obj.model_call_details["litellm_params"] = { - "metadata": { - "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" - } + "metadata": {"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"} } logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] @@ -513,16 +454,11 @@ class TestAzureAnthropicCostCalculation: ) mock_completion_cost.assert_called_once() - assert ( - mock_completion_cost.call_args[1]["model"] - == "anthropic/claude-3-5-haiku-20241022" - ) + assert mock_completion_cost.call_args[1]["model"] == "anthropic/claude-3-5-haiku-20241022" assert kwargs["response_cost"] == 0.002 @patch("litellm.completion_cost") - def test_cost_calculation_skips_unknown_litellm_params_model_for_model_group( - self, mock_completion_cost - ): + def test_cost_calculation_skips_unknown_litellm_params_model_for_model_group(self, mock_completion_cost): """When litellm_params.model is itself the "unknown" sentinel, the deployment-model branch must not short-circuit; resolution falls through to model_group so costing still prices the real model instead of "unknown".""" @@ -535,9 +471,7 @@ class TestAzureAnthropicCostCalculation: logging_obj = self._create_mock_logging_obj(model="unknown") logging_obj.model_call_details["litellm_params"] = { "model": "unknown", - "metadata": { - "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" - }, + "metadata": {"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"}, } logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] @@ -555,17 +489,12 @@ class TestAzureAnthropicCostCalculation: ) mock_completion_cost.assert_called_once() - assert ( - mock_completion_cost.call_args[1]["model"] - == "anthropic/claude-3-5-haiku-20241022" - ) + assert mock_completion_cost.call_args[1]["model"] == "anthropic/claude-3-5-haiku-20241022" assert kwargs["response_cost"] == 0.003 assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022" @patch("litellm.completion_cost") - def test_streaming_cost_calculation_resolves_model_from_message_start_chunk( - self, mock_completion_cost - ): + def test_streaming_cost_calculation_resolves_model_from_message_start_chunk(self, mock_completion_cost): """On the bare /anthropic passthrough path litellm_params carries no model or model_group and the body model is the "unknown" sentinel; the model must be recovered from the message_start SSE event so completion_cost @@ -628,9 +557,7 @@ class TestAzureAnthropicCostCalculation: ), _sse("message_stop", {"type": "message_stop"}), ] - all_chunks = list( - PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames) - ) + all_chunks = list(PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames)) logging_obj = RealLoggingObj( model="unknown", @@ -674,9 +601,7 @@ class TestAzureAnthropicCostCalculation: ] assert ( - AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( - chunks - ) + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(chunks) == "claude-3-5-haiku-20241022" ) @@ -692,14 +617,11 @@ class TestAzureAnthropicCostCalculation: ) assert ( - AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( - [raw_event] - ) + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks([raw_event]) == "claude-3-5-haiku-20241022" ) - class TestAnthropicBatchPassthroughCostTracking: """Test cases for Anthropic batch passthrough cost tracking functionality""" @@ -799,9 +721,7 @@ class TestAnthropicBatchPassthroughCostTracking: ) mock_batches_config_instance = MagicMock() - mock_batches_config_instance.transform_retrieve_batch_response.return_value = ( - mock_batch_response - ) + mock_batches_config_instance.transform_retrieve_batch_response.return_value = mock_batch_response mock_batches_config.return_value = mock_batches_config_instance # Test the handler @@ -929,9 +849,7 @@ class TestAnthropicBatchPassthroughCostTracking: "transform_retrieve_batch_response", return_value=mock_batch_response, ): - with patch.object( - AnthropicPassthroughLoggingHandler, "_store_batch_managed_object" - ): + with patch.object(AnthropicPassthroughLoggingHandler, "_store_batch_managed_object"): result = AnthropicPassthroughLoggingHandler.batch_creation_handler( httpx_response=mock_httpx_response, logging_obj=mock_logging_obj, @@ -946,10 +864,7 @@ class TestAnthropicBatchPassthroughCostTracking: # Verify unified_object_id contains anthropic/ prefix unified_object_id = result["kwargs"]["unified_object_id"] decoded = base64.urlsafe_b64decode(unified_object_id + "==").decode() - assert ( - "anthropic/claude-sonnet-4-5-20250929" in decoded - or "claude-sonnet-4-5-20250929" in decoded - ) + assert "anthropic/claude-sonnet-4-5-20250929" in decoded or "claude-sonnet-4-5-20250929" in decoded @pytest.mark.parametrize( "kwargs,expected_user_id,expected_team_id", @@ -1007,9 +922,7 @@ class TestAnthropicBatchPassthroughCostTracking: mock_managed_files_hook.store_unified_object_id = AsyncMock() with ( patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl, - patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger" - ), + patch("litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger"), ): mock_pl.get_proxy_hook.return_value = mock_managed_files_hook AnthropicPassthroughLoggingHandler._store_batch_managed_object( @@ -1060,9 +973,7 @@ class TestAnthropicBatchPassthroughCostTracking: is never tracked. The failure has to reach the log instead of being reported as a success.""" mock_managed_files_hook = MagicMock() - mock_managed_files_hook.store_unified_object_id = AsyncMock( - side_effect=RuntimeError("db down") - ) + mock_managed_files_hook.store_unified_object_id = AsyncMock(side_effect=RuntimeError("db down")) with ( patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl, patch( @@ -1104,9 +1015,7 @@ class TestAnthropicBatchPassthroughCostTracking: this gate an id-scoped route reached the store with a mismatched id, where it could only either claim a row it did not create or fail the model_object_id unique constraint.""" - with patch.object( - AnthropicPassthroughLoggingHandler, "_store_batch_managed_object" - ) as mock_store: + with patch.object(AnthropicPassthroughLoggingHandler, "_store_batch_managed_object") as mock_store: AnthropicPassthroughLoggingHandler.batch_creation_handler( httpx_response=mock_httpx_response, logging_obj=mock_logging_obj, @@ -1120,9 +1029,7 @@ class TestAnthropicBatchPassthroughCostTracking: assert mock_store.call_count == (1 if registers else 0) - def test_batch_creation_handler_failure_status_code( - self, mock_logging_obj, mock_request_body - ): + def test_batch_creation_handler_failure_status_code(self, mock_logging_obj, mock_request_body): """Test batch creation handler with non-200 status code""" mock_response = MagicMock() mock_response.status_code = 400 @@ -1145,9 +1052,7 @@ class TestAnthropicBatchPassthroughCostTracking: assert result["kwargs"]["response_cost"] == 0.0 @patch("litellm.proxy.proxy_server.proxy_logging_obj") - def test_store_batch_managed_object_success( - self, mock_proxy_logging_obj, mock_logging_obj - ): + def test_store_batch_managed_object_success(self, mock_proxy_logging_obj, mock_logging_obj): """Test storing batch managed object""" from litellm.types.utils import LiteLLMBatch @@ -1178,9 +1083,7 @@ class TestAnthropicBatchPassthroughCostTracking: ) # Verify managed files hook was called - mock_proxy_logging_obj.get_proxy_hook.assert_called_once_with( - "managed_files" - ) + mock_proxy_logging_obj.get_proxy_hook.assert_called_once_with("managed_files") class TestBuildCompleteStreamingResponseRobustness: @@ -1343,11 +1246,7 @@ class TestPureTextFastPathParity: }, ) ) - frames.append( - self._sse( - "content_block_stop", {"type": "content_block_stop", "index": b} - ) - ) + frames.append(self._sse("content_block_stop", {"type": "content_block_stop", "index": b})) frames.append( self._sse( "message_delta", @@ -1405,9 +1304,7 @@ class TestPureTextFastPathParity: # And the full logged payload (sans non-deterministic response id). def _scrub(p): d = dict(p) - r = d.get("complete_streaming_response_in_db") or d.get( - "complete_streaming_response" - ) + r = d.get("complete_streaming_response_in_db") or d.get("complete_streaming_response") return d, getattr(r, "usage", None) assert _scrub(k_fast)[1] == _scrub(k_legacy)[1] @@ -1419,11 +1316,7 @@ class TestPureTextFastPathParity: self._assert_parity(self._text_stream(["Just one piece of text."])) def test_parity_cache_tokens(self): - self._assert_parity( - self._text_stream( - ["a", "b", "c"], input_tokens=20, cache_creation=5, cache_read=7 - ) - ) + self._assert_parity(self._text_stream(["a", "b", "c"], input_tokens=20, cache_creation=5, cache_read=7)) def test_parity_max_tokens_stop(self): self._assert_parity(self._text_stream(["tok"] * 8, stop_reason="max_tokens")) @@ -1484,12 +1377,7 @@ class TestPureTextFastPathParity: self._sse("message_stop", {"type": "message_stop"}), ] all_chunks = self._to_all_chunks(frames) - assert ( - AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks( - list(all_chunks) - ) - is None - ) + assert AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(list(all_chunks)) is None def test_collapse_returns_none_for_thinking(self): frames = [ @@ -1526,19 +1414,12 @@ class TestPureTextFastPathParity: self._sse("message_stop", {"type": "message_stop"}), ] all_chunks = self._to_all_chunks(frames) - assert ( - AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks( - list(all_chunks) - ) - is None - ) + assert AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(list(all_chunks)) is None def test_collapse_actually_shrinks_chunk_count(self): frames = self._text_stream(["a"] * 50) all_chunks = list(self._to_all_chunks(frames)) - collapsed = AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks( - all_chunks - ) + collapsed = AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) assert collapsed is not None # 50 text deltas + 50 event markers + 1 ping collapse to far fewer. assert len(collapsed) < len(all_chunks) / 2 @@ -1604,10 +1485,7 @@ class TestPureTextFastPathParity: self._sse("message_stop", {"type": "message_stop"}), ] all_chunks = list(self._to_all_chunks(frames)) - assert ( - AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) - is None - ) + assert AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) is None class TestInterruptedStreamOutputTokenRecovery: @@ -1722,9 +1600,7 @@ class TestInterruptedStreamOutputTokenRecovery: import litellm placeholder = 2 - result = self._run( - self._interrupted_chunks(placeholder_output_tokens=placeholder) - ) + result = self._run(self._interrupted_chunks(placeholder_output_tokens=placeholder)) usage = result["result"].usage expected = litellm.token_counter( @@ -1897,9 +1773,7 @@ class TestStreamFalseDeduplication: "delta": {"type": "text_delta", "text": "Hello"}, }, ), - TestStreamFalseDeduplication._sse( - "content_block_stop", {"type": "content_block_stop", "index": 0} - ), + TestStreamFalseDeduplication._sse("content_block_stop", {"type": "content_block_stop", "index": 0}), TestStreamFalseDeduplication._sse( "message_delta", { @@ -1949,15 +1823,12 @@ class TestStreamFalseDeduplication: # The assembled response must be stored on model_call_details so callbacks # can identify this as a completed streaming call, not an in-progress one. - assert ( - logging_obj.model_call_details.get("complete_streaming_response") - is not None - ), "complete_streaming_response must be set on model_call_details after assembly" + assert logging_obj.model_call_details.get("complete_streaming_response") is not None, ( + "complete_streaming_response must be set on model_call_details after assembly" + ) # The returned result must match what was stored - assert result["result"] is logging_obj.model_call_details.get( - "complete_streaming_response" - ) + assert result["result"] is logging_obj.model_call_details.get("complete_streaming_response") def test_dedup_guard_fires_when_stream_true_on_logging_obj(self): """ @@ -1986,21 +1857,16 @@ class TestStreamFalseDeduplication: assert logging_obj._is_assembled_stream_success(result=mock_response) is True # First dispatch sets the flag - assert not logging_obj.model_call_details.get( - "has_dispatched_final_stream_success" - ) + assert not logging_obj.model_call_details.get("has_dispatched_final_stream_success") logging_obj.model_call_details["has_dispatched_final_stream_success"] = True # Second dispatch would be blocked — simulate the guard check would_skip = bool( logging_obj._is_assembled_stream_success(result=mock_response) - and logging_obj.model_call_details.get( - "has_dispatched_final_stream_success" - ) + and logging_obj.model_call_details.get("has_dispatched_final_stream_success") ) assert would_skip is True, ( - "Dedup guard must block a second dispatch_success_handlers call for the " - "same assembled streaming response" + "Dedup guard must block a second dispatch_success_handlers call for the same assembled streaming response" ) def test_sse_fallback_path_sets_stream_true_for_dedup(self): @@ -2034,9 +1900,7 @@ class TestStreamFalseDeduplication: would_skip = bool( logging_obj._is_assembled_stream_success(result=mock_response) - and logging_obj.model_call_details.get( - "has_dispatched_final_stream_success" - ) + and logging_obj.model_call_details.get("has_dispatched_final_stream_success") ) assert would_skip is True @@ -2100,9 +1964,9 @@ class TestNonStreamingResponseRedaction: logging_obj=logging_obj, ) - assert ( - "complete_streaming_response" not in logging_obj.model_call_details - ), "non-streaming responses must not populate complete_streaming_response" + assert "complete_streaming_response" not in logging_obj.model_call_details, ( + "non-streaming responses must not populate complete_streaming_response" + ) def test_streaming_sets_complete_streaming_response(self): from litellm.types.utils import ModelResponse @@ -2119,10 +1983,7 @@ class TestNonStreamingResponseRedaction: logging_obj=logging_obj, ) - assert ( - logging_obj.model_call_details.get("complete_streaming_response") - is response - ) + assert logging_obj.model_call_details.get("complete_streaming_response") is response def test_non_streaming_response_is_redacted_when_message_logging_off(self): from litellm.litellm_core_utils.redact_messages import ( @@ -2196,10 +2057,8 @@ class TestAnthropicUsageOnlyFallback: ] def test_build_usage_only_recovers_cache_inclusive_usage(self): - response = ( - AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( - all_chunks=self._CHUNKS, model="claude-3-5-haiku-20241022" - ) + response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=self._CHUNKS, model="claude-3-5-haiku-20241022" ) assert response is not None usage = response.usage @@ -2259,10 +2118,8 @@ class TestAnthropicUsageOnlyFallback: } ), ] - response = ( - AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( - all_chunks=chunks, model="unknown" - ) + response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=chunks, model="unknown" ) assert response is not None assert response.model == "claude-opus-4-6" @@ -2287,9 +2144,7 @@ class TestAnthropicUsageOnlyFallback: ('data: {"a": 1}', {"a": 1}), ], ) - def test_extract_sse_data_handles_malformed_and_sentinel_lines( - self, event_str, expected - ): + def test_extract_sse_data_handles_malformed_and_sentinel_lines(self, event_str, expected): assert extract_sse_data(event_str) == expected def _real_logging_obj(self): @@ -2309,12 +2164,8 @@ class TestAnthropicUsageOnlyFallback: return logging_obj @patch("litellm.completion_cost") - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) - def test_handler_falls_back_when_assembly_returns_none( - self, mock_assemble, mock_cost - ): + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") + def test_handler_falls_back_when_assembly_returns_none(self, mock_assemble, mock_cost): mock_assemble.return_value = None mock_cost.return_value = 0.0021 logging_obj = self._real_logging_obj() @@ -2335,9 +2186,7 @@ class TestAnthropicUsageOnlyFallback: assert result["kwargs"]["response_cost"] == 0.0021 @patch("litellm.completion_cost") - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") def test_handler_falls_back_when_assembly_raises(self, mock_assemble, mock_cost): import litellm @@ -2367,9 +2216,7 @@ class TestAnthropicUsageOnlyFallback: assert result["result"].usage.completion_tokens == 55 assert result["kwargs"]["response_cost"] == 0.0021 - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") def test_handler_returns_none_when_no_usage_recoverable(self, mock_assemble): # assembly fails AND the chunks carry no usage event, so there is nothing # to price; the handler must return None rather than fabricate a response @@ -2391,15 +2238,9 @@ class TestAnthropicUsageOnlyFallback: assert result["result"] is None assert result["kwargs"] == {} - @patch.object( - AnthropicPassthroughLoggingHandler, "build_usage_only_response_from_chunks" - ) - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) - def test_handler_does_not_crash_when_usage_only_fallback_raises( - self, mock_assemble, mock_fallback - ): + @patch.object(AnthropicPassthroughLoggingHandler, "build_usage_only_response_from_chunks") + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") + def test_handler_does_not_crash_when_usage_only_fallback_raises(self, mock_assemble, mock_fallback): # if the usage-only fallback itself raises, it must be treated as None and # drop gracefully, not propagate out and crash the success handler mock_assemble.return_value = None @@ -2427,7 +2268,6 @@ class TestAnthropicResponseCostRecordedOnModelCallDetails: builder must record it there or streaming pass-through logs $0.""" - class TestAnthropicPassthroughFastMode: """Anthropic charges a provider-specific multiplier for ``speed=fast``, applied off ``usage.speed`` and covering every token type, cache included. The response usage @@ -2436,20 +2276,20 @@ class TestAnthropicPassthroughFastMode: MODEL = "claude-opus-4-8" STREAM_CHUNKS = [ - 'event: message_start', + "event: message_start", 'data: {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant",' ' "model": "claude-opus-4-8", "content": [], "stop_reason": null,' ' "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 0}}}', - 'event: content_block_start', + "event: content_block_start", 'data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}', - 'event: content_block_delta', + "event: content_block_delta", 'data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}', - 'event: content_block_stop', + "event: content_block_stop", 'data: {"type": "content_block_stop", "index": 0}', - 'event: message_delta', + "event: message_delta", 'data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"},' ' "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 100}}', - 'event: message_stop', + "event: message_stop", 'data: {"type": "message_stop"}', ] @@ -2643,58 +2483,6 @@ class TestRecordPartialUsageForFailure: assert "response_cost" not in logging_obj.model_call_details -class _FailureRecorder(CustomLogger): - def __init__(self): - super().__init__() - self.failure_kwargs = [] - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - self.failure_kwargs.append(kwargs) - - -class TestLogErrorFrameAsFailure: - """A stream closed at the provider's `event: error` frame is one failed attempt: its failure log bills the - usage streamed before the frame on the attempt-scoped copy its caller hands it, and the request's shared - logging state stays as it was, so the attempt the router opens next logs with nothing left behind.""" - - @pytest.mark.asyncio - async def test_logs_a_failure_with_the_usage_streamed_before_the_frame_and_leaves_the_request_state_alone(self): - recorder = _FailureRecorder() - logging_obj = LiteLLMLoggingObj( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "hello"}], - stream=True, - call_type="anthropic_messages", - start_time=datetime.now(), - litellm_call_id="test-error-frame-failure", - function_id="test-error-frame-failure", - dynamic_async_failure_callbacks=[recorder], - ) - chunks = [ - *TestRecordPartialUsageForFailure()._interrupted_chunks()[:1], - TestRecordPartialUsageForFailure._sse( - "error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}} - ), - ] - - await AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( - attempt_logging_obj=logging_obj.attempt_scoped_copy(), - request_body={"model": "claude-sonnet-5", "stream": True}, - all_chunks=chunks, - error_event=("overloaded_error", "Overloaded", 529), - ) - - assert len(recorder.failure_kwargs) == 1 - failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] - assert failure_payload["status"] == "failure" - assert failure_payload["prompt_tokens"] == 52 - assert failure_payload["response_cost"] > 0 - assert isinstance(recorder.failure_kwargs[0]["exception"], litellm.InternalServerError) - assert "combined_usage_object" not in logging_obj.model_call_details - assert "response_cost" not in logging_obj.model_call_details - assert logging_obj.should_run_logging("async_failure") - - @pytest.mark.parametrize( ("all_chunks", "interrupted"), [ @@ -2715,9 +2503,7 @@ def test_stream_was_interrupted_skips_data_lines_that_are_not_json_objects( @pytest.mark.asyncio -async def test_anthropic_passthrough_handler( - mock_httpx_response, mock_response, mock_logging_obj -): +async def test_anthropic_passthrough_handler(mock_httpx_response, mock_response, mock_logging_obj): """ Unit test - Assert that the anthropic passthrough handler calls the litellm logging object's async_success_handler """ @@ -3058,3 +2844,102 @@ def test_streamed_messages_end_turn_logs_no_stop_details(mock_logging_obj): assert logged.choices[0].finish_reason == "stop" provider_fields: Final = logged.choices[0].message.provider_specific_fields or {} assert "stop_details" not in provider_fields + + +def _go_sse(event: str, data: dict) -> str: + return f"event: {event}\ndata: {json.dumps(data)}\n\n" + + +_GO_MESSAGE_DELTA_USAGE: Final = { + "input_tokens": 1200, + "output_tokens": 300, + "cost": 0, + "opencode": {"billing_source": "go", "list_cost": 0.00412}, +} + + +def _go_stream_chunks() -> list[str]: + message = { + "id": "msg_go", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 1200, "output_tokens": 0}, + } + return [ + _go_sse("message_start", {"type": "message_start", "message": message}), + _go_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _go_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + ), + _go_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _go_sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": _GO_MESSAGE_DELTA_USAGE, + }, + ), + _go_sse("message_stop", {"type": "message_stop"}), + ] + + +def _go_logging_obj() -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model="opencode_go/claude-haiku-5-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="go-stream", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "opencode_go" + logging_obj.model_call_details["stream"] = True + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + return logging_obj + + +def test_streamed_messages_bill_the_list_cost_reported_in_the_final_message_delta(): + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=_go_logging_obj(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "opencode_go/claude-haiku-5-5", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=_go_stream_chunks(), + end_time=datetime.now(), + ) + + assert result["kwargs"]["response_cost"] == 0.00412 + assert result["result"]._hidden_params["response_cost"] == 0.00412 + + +def test_streamed_messages_from_other_providers_ignore_a_zero_usage_cost_field(): + logging_obj = _go_logging_obj() + logging_obj.model_call_details["custom_llm_provider"] = "anthropic" + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "claude-haiku-5-5", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=_go_stream_chunks(), + end_time=datetime.now(), + ) + + headers = result["result"]._hidden_params.get("additional_headers") or {} + assert "llm_provider-x-litellm-response-cost" not in headers + assert result["kwargs"].get("response_cost") != 0.0 diff --git a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py new file mode 100644 index 00000000000..354c65776dc --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py @@ -0,0 +1,185 @@ +import asyncio +import json +from collections.abc import Iterator +from datetime import datetime +from typing import Final + +import httpx +import pytest +import respx +from fastapi.testclient import TestClient + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy import proxy_server +from tests._master_key import MASTER_KEY + +_BASE: Final = "https://chatgpt.com/backend-api/codex" +_TOKEN: Final = "test-subscription-token" +_HEADERS: Final = {"authorization": f"Bearer {_TOKEN}", "x-litellm-api-key": MASTER_KEY} + + +@pytest.fixture +def client(monkeypatch: pytest.MonkeyPatch, httpx_transport: None) -> Iterator[TestClient]: + monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY) + yield TestClient(proxy_server.app) + + +def test_native_catalog_preserves_account_and_filters_gateway_credentials( + client: TestClient, respx_mock: respx.MockRouter +) -> None: + catalog: Final = {"models": [{"slug": "subscription-model", "supported_reasoning_levels": ["high"]}]} + upstream: Final = respx_mock.get(f"{_BASE}/models", params={"client_version": "1.0.0"}).respond(json=catalog) + response: Final = client.get( + "/backend-api/codex/models?client_version=1.0.0", + headers={ + **_HEADERS, + "chatgpt-account-id": "account-a", + "cookie": "session=test-cookie", + "x-pass-authorization": "Bearer test-injected-token", + "x-pass-x-litellm-api-key": "test-injected-key", + }, + ) + assert response.status_code == 200 + assert response.json() == catalog + assert upstream.calls.last.request.headers["authorization"] == f"Bearer {_TOKEN}" + assert upstream.calls.last.request.headers["chatgpt-account-id"] == "account-a" + assert "x-litellm-api-key" not in upstream.calls.last.request.headers + assert "cookie" not in upstream.calls.last.request.headers + + +@pytest.mark.parametrize("account", ["account-a", "account-b"]) +def test_native_responses_preserves_body_and_stream( + client: TestClient, respx_mock: respx.MockRouter, account: str +) -> None: + body: Final = { + "model": "subscription-model", + "input": [{"role": "user", "content": "Hello"}], + "reasoning": {"effort": "high"}, + "store": False, + "stream": True, + } + stream: Final = ( + b'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"Hello"}\n\n' + ) + upstream: Final = respx_mock.post(f"{_BASE}/responses").respond( + content=stream, headers={"content-type": "text/event-stream"} + ) + response: Final = client.post( + "/backend-api/codex/responses", + json=body, + headers={**_HEADERS, "authorization": f"Bearer test-{account}", "chatgpt-account-id": account}, + ) + assert response.status_code == 200 + assert response.content == stream + assert response.headers["content-type"].startswith("text/event-stream") + assert json.loads(upstream.calls.last.request.content) == body + assert upstream.calls.last.request.headers["authorization"] == f"Bearer test-{account}" + assert upstream.calls.last.request.headers["chatgpt-account-id"] == account + + +@pytest.mark.parametrize("gateway_key", [None, "sk-invalid-test-key"]) +def test_gateway_auth_rejects_before_upstream( + client: TestClient, respx_mock: respx.MockRouter, gateway_key: str | None +) -> None: + headers: Final = { + "authorization": f"Bearer {_TOKEN}", + **({"x-litellm-api-key": gateway_key} if gateway_key else {}), + } + response: Final = client.get("/backend-api/codex/models", headers=headers) + assert response.status_code in (400, 401, 403) + assert not respx_mock.calls + + +@pytest.mark.parametrize("authorization", ["", "Basic test", "Bearer ", f"Bearer {MASTER_KEY}"]) +def test_subscription_credential_required(client: TestClient, respx_mock: respx.MockRouter, authorization: str) -> None: + response: Final = client.get("/backend-api/codex/models", headers={**_HEADERS, "authorization": authorization}) + assert response.status_code == 401 + assert not respx_mock.calls + + +@pytest.mark.parametrize("status", [401, 403, 429]) +def test_provider_errors_do_not_fall_back(client: TestClient, respx_mock: respx.MockRouter, status: int) -> None: + upstream: Final = respx_mock.post(f"{_BASE}/responses").respond( + status, json={"error": {"message": "subscription unavailable"}}, headers={"retry-after": "60"} + ) + response: Final = client.post( + "/backend-api/codex/responses", headers=_HEADERS, json={"model": "subscription-model", "stream": True} + ) + assert response.status_code == status + assert upstream.call_count == 1 + + +@pytest.mark.parametrize("method,path", [("POST", "models"), ("GET", "responses"), ("GET", "deviceauth")]) +def test_only_native_inference_routes_are_forwarded( + client: TestClient, respx_mock: respx.MockRouter, method: str, path: str +) -> None: + response: Final = client.request(method, f"/backend-api/codex/{path}", headers=_HEADERS) + assert response.status_code in (405, 422) + assert not respx_mock.calls + + +class UsageRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + await self.events.put(kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_usage_logged_without_credentials( + monkeypatch: pytest.MonkeyPatch, httpx_transport: None, respx_mock: respx.MockRouter, stream: bool +) -> None: + monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY) + recorder: Final = UsageRecorder() + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + litellm.register_model( + { + "chatgpt/subscription-model": { + "litellm_provider": "chatgpt", + "mode": "responses", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + } + ) + result: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "subscription-model", + "output": [], + "usage": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15}, + } + event: Final = {"type": "response.completed", "response": result} + response_body: Final = ( + httpx.Response( + 200, + content=f"event: response.completed\ndata: {json.dumps(event)}\n\n", + headers={"content-type": "text/event-stream"}, + ) + if stream + else httpx.Response(200, json=result) + ) + respx_mock.post(f"{_BASE}/responses").mock(return_value=response_body) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test") as caller: + response: Final = await caller.post( + "/backend-api/codex/responses", + json={"model": "subscription-model", "stream": stream, "input": "Hello"}, + headers={**_HEADERS, "x-pass-authorization": "Bearer test-hidden-token", "cookie": "test-hidden-cookie"}, + ) + assert response.status_code == 200 + logged: Final = await asyncio.wait_for(recorder.events.get(), timeout=5) + payload: Final = logged["standard_logging_object"] + assert isinstance(payload, dict) + assert payload["prompt_tokens"] == 12 + assert payload["completion_tokens"] == 3 + serialized: Final = json.dumps(logged, default=str) + for secret in (_TOKEN, MASTER_KEY, "test-hidden-token", "test-hidden-cookie"): + assert secret not in serialized 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 0186a33dbc4..60043daa402 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 @@ -4357,6 +4357,7 @@ class TestVertexRawPredictStreamingClassification: "type": "http", "method": "POST", "path": f"/vertex_ai/{endpoint}", + "path_params": {"endpoint": endpoint}, "headers": [ (b"content-type", b"application/json"), (b"x-litellm-api-key", b"test-key"), diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f0baceecdaa..db0a74d360a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -8536,6 +8536,7 @@ PROXY_STAMPED_NAMES: Final = frozenset( "litellm_trusted_callback_vars", "_litellm_addressed_response_id", "_litellm_strip_stream_usage", + "_litellm_stream_usage_injected", "client_side_timeout", "model_file_id_mapping", ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py index da084146894..7d56487a844 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py @@ -220,7 +220,11 @@ async def test_failed_anthropic_stream_records_partial_usage_off_the_event_loop( logging_obj.model_call_details = {"model": model, "stream": True} logging_obj.litellm_params = {} logging_obj.get_router_model_id.return_value = None + logging_obj.attempt_scoped_copy.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() + logging_obj.dispatch_attempt_failure_handlers = AsyncMock( + side_effect=logging_obj.dispatch_failure_handlers + ) _, took, lags = await timed_with_loop_lags( lambda: PassThroughStreamingHandler.schedule_stream_failure_logging( diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 8ce41129aec..4aa7e9cfbae 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -2,7 +2,9 @@ import asyncio import json +from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -22,6 +24,8 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import aclo from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) +from litellm.proxy.utils import _failure_fields_to_lift +from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType @@ -407,7 +411,7 @@ async def _collect_openai_passthrough_chunks(chunks, endpoint_type): request_body={"model": "gpt-4o-mini", "stream": True}, litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=endpoint_type, - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), passthrough_success_handler_obj=MagicMock(), url_route="/openai/v1/chat/completions", route_streaming_logging=AsyncMock(), @@ -705,7 +709,7 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( messages=[{"role": "user", "content": "hi"}], stream=True, call_type="anthropic_messages", - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), litellm_call_id="test-mid-stream-timeout", function_id="test-mid-stream-timeout", dynamic_async_success_callbacks=[recorder], @@ -750,6 +754,109 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) +def _anthropic_stream_dropping_after(frames: tuple[bytes, ...]) -> httpx.Response: + mock: Final = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes() -> AsyncIterator[bytes]: + for frame in frames: + yield frame + raise httpx.ReadError("upstream closed the connection") + + mock.aiter_bytes = _aiter_bytes + return mock + + +def _anthropic_messages_attempt( + logging_obj: LiteLLMLoggingObj, response: httpx.Response +) -> AsyncGenerator[bytes, None]: + return PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-sonnet-5", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime(2026, 1, 1), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + + +def _anthropic_messages_logging_obj(recorder: _EventRecorder, call_id: str) -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + + +async def _drop_after_message_start(logging_obj: LiteLLMLoggingObj, message_id: str, input_tokens: int) -> None: + dropped: Final = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_dropping_after((_anthropic_message_start_frame(message_id, input_tokens),)) + ) + with pytest.raises(httpx.ReadError): + async for _ in dropped: + pass + + +@pytest.mark.asyncio +async def test_dropped_attempt_logs_its_failure_without_claiming_the_shared_logging_object() -> None: + """The router retries a /v1/messages stream the provider drops after message_start on the same logging + object. The dropped attempt's failure log bills its own usage and leaves the shared object free: the + retry's success log carries the retry's usage and no stale exception, while the stash stays for the + failure row of a request that runs out of retries until the router opens the next attempt.""" + recorder: Final = _EventRecorder() + logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-dropped-attempt") + + await _drop_after_message_start(logging_obj, "msg_a1", 52) + await _settle(recorder.failure_kwargs) + + assert recorder.failure_kwargs[0]["standard_logging_object"]["prompt_tokens"] == 52 + assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadError) + assert logging_obj.model_call_details.get("has_logged_async_failure") is not True + assert "exception" not in logging_obj.model_call_details + stashed: Final = logging_obj.model_call_details.get("combined_usage_object") + assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 52 + + logging_obj.carry_failed_attempt_usage() + retry: Final = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) + ) + async for _ in retry: + pass + await _settle(recorder.success_kwargs) + + success_payload: Final = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) + assert recorder.success_kwargs[0].get("exception") is None + assert len(recorder.failure_kwargs) == 1 + + +@pytest.mark.asyncio +async def test_every_dropped_attempt_logs_its_own_failure() -> None: + """Two attempts of one /v1/messages request dropped in a row each log a failure billing their own usage; + the first attempt's log must not swallow the second's, and the shared object ends up holding the last + attempt's usage for the request's failure row.""" + recorder: Final = _EventRecorder() + logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-every-dropped-attempt") + + await _drop_after_message_start(logging_obj, "msg_a1", 52) + await _settle(recorder.failure_kwargs) + logging_obj.carry_failed_attempt_usage() + await _drop_after_message_start(logging_obj, "msg_a2", 9) + await _settle(recorder.failure_kwargs, count=2) + + assert [kwargs["standard_logging_object"]["prompt_tokens"] for kwargs in recorder.failure_kwargs] == [52, 9] + assert all(isinstance(kwargs["exception"], httpx.ReadError) for kwargs in recorder.failure_kwargs) + stashed: Final = logging_obj.model_call_details.get("combined_usage_object") + assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 9 + def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes: return _anthropic_sse( "message_start", @@ -788,9 +895,9 @@ def _anthropic_completed_stream_tail() -> tuple[bytes, ...]: ) -async def _settle(recorded: list[dict[str, object]]) -> None: +async def _settle(recorded: Sequence[Mapping[str, object]], count: int = 1) -> None: for _ in range(300): - if recorded: + if len(recorded) >= count: return await asyncio.sleep(0.01) @@ -857,8 +964,6 @@ async def test_chunk_processor_logs_a_provider_error_frame_as_a_failure_when_clo assert received == list(frames) assert recorder.success_kwargs == [] assert logging_obj.model_call_details.get("has_dispatched_final_stream_success") is not True - assert "combined_usage_object" not in logging_obj.model_call_details - assert "response_cost" not in logging_obj.model_call_details assert logging_obj.should_run_logging("async_failure") assert len(recorder.failure_kwargs) == 1 failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] @@ -991,6 +1096,31 @@ async def test_superseded_attempt_failure_log_keeps_its_own_deployment_when_the_ assert logging_obj.model_call_details["litellm_params"]["model_info"]["id"] == "deployment-2" +@pytest.mark.asyncio +async def test_error_frame_then_exhausted_retry_bills_both_attempts() -> None: + recorder: Final = _EventRecorder() + logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-error-frame-terminal-billing") + superseded: Final = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())) + ) + async for chunk in superseded: + if b"event: error" in chunk: + break + await aclose_superseded_attempt(superseded) + carry_superseded_attempt_usage({"litellm_logging_obj": logging_obj}) + await _drop_after_message_start(logging_obj, "msg_a2", 9) + await _settle(recorder.failure_kwargs, count=2) + + failures: Final = tuple(event["standard_logging_object"] for event in recorder.failure_kwargs) + assert tuple(event["prompt_tokens"] for event in failures) == (52, 9) + terminal: Final = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}) + usage: Final = terminal["combined_usage_object"] + assert isinstance(usage, litellm.Usage) + assert usage.prompt_tokens == sum(event["prompt_tokens"] for event in failures) + assert usage.completion_tokens == sum(event["completion_tokens"] for event in failures) + assert terminal["response_cost"] == pytest.approx(sum(event["response_cost"] for event in failures)) + + @pytest.mark.asyncio async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_frame(): """Both attempts of a retried /v1/messages stream log through the request's one logging object. The @@ -1025,6 +1155,7 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr if b"event: error" in chunk: break await aclose_superseded_attempt(superseded) + carry_superseded_attempt_usage({"litellm_logging_obj": logging_obj}) await _settle(recorder.failure_kwargs) retry = attempt((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) async for _ in retry: @@ -1036,8 +1167,6 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr success_payload = recorder.success_kwargs[0]["standard_logging_object"] assert success_payload["status"] == "success" assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) - - def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], diff --git a/tests/unit/proxy/proxy_server/test_openapi_customization.py b/tests/unit/proxy/proxy_server/test_openapi_customization.py index 6c588976efc..0eacf333960 100644 --- a/tests/unit/proxy/proxy_server/test_openapi_customization.py +++ b/tests/unit/proxy/proxy_server/test_openapi_customization.py @@ -36,6 +36,7 @@ from litellm.proxy.proxy_server import ( get_openapi_schema, mount_swagger_ui, ) +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder from .conftest import normalize @@ -209,6 +210,17 @@ def test_inject_websocket_stubs_into_openapi_schema_adds_stub(): } +def test_inject_websocket_stubs_into_openapi_schema_documents_query_params_that_forward_their_fields(): + query_param: Final = AttributeForwarder(SimpleNamespace(name="session_id", required=True)) + route: Final = SimpleNamespace(path="/ws/chat", name="ws_chat", dependant=SimpleNamespace(query_params=[query_param])) + + result: Final = _inject_websocket_stubs_into_openapi_schema({"paths": {}}, [route]) + + assert result["paths"]["/ws/chat"]["get"]["parameters"] == [ + {"name": "session_id", "in": "query", "required": True, "schema": {"type": "string"}} + ] + + def test_inject_websocket_stubs_into_openapi_schema_does_not_overwrite_existing_get(): # Existing GET on the same path must not be replaced by the stub. existing_get = {"summary": "real http get", "operationId": "real_get"} diff --git a/tests/unit/proxy/proxy_server/test_routes_config.py b/tests/unit/proxy/proxy_server/test_routes_config.py index a92986392e9..7d0a7f3f261 100644 --- a/tests/unit/proxy/proxy_server/test_routes_config.py +++ b/tests/unit/proxy/proxy_server/test_routes_config.py @@ -716,6 +716,53 @@ def test_config_list_exposes_config_reload_interval(client, auth_as, mock_prisma assert entry["field_default_value"] == 30 +@pytest.mark.parametrize("enabled", (True, False)) +def test_config_list_exposes_and_saves_wildcard_routes_toggle( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[object]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, + enabled: bool, +) -> None: + from pydantic import TypeAdapter + + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import ConfigList, LitellmUserRoles + + table: Final = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + _seed_settings_store(monkeypatch, {}) + monkeypatch.setattr(ps, "general_settings", ps.proxy_config.settings) + entries: Final = TypeAdapter(list[ConfigList]) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + before: Final = client.get("/config/list", params={"config_type": "general_settings"}) + assert before.status_code == 200, before.text + toggle: Final = next( + entry for entry in entries.validate_json(before.content) if entry.field_name == "return_wildcard_routes" + ) + assert toggle.field_type == "Boolean" + assert toggle.field_value is None + assert toggle.editable is True + + saved: Final = client.post( + "/config/field/update", + json={"field_name": "return_wildcard_routes", "field_value": enabled, "config_type": "general_settings"}, + ) + assert saved.status_code == 200, saved.text + persisted: Final = TypeAdapter(dict[str, bool]).validate_json( + table.upsert.call_args.kwargs["data"]["create"]["param_value"] + ) + assert persisted == {"return_wildcard_routes": enabled} + + after: Final = client.get("/config/list", params={"config_type": "general_settings"}) + assert after.status_code == 200, after.text + updated: Final = next( + entry for entry in entries.validate_json(after.content) if entry.field_name == "return_wildcard_routes" + ) + assert updated.field_value is enabled + + def test_config_field_update_accepts_config_reload_interval(client, auth_as, mock_prisma, monkeypatch): """POST /config/field/update accepts proxy_config_reload_interval_seconds and persists it to the DB general_settings row for all pods to pick up.""" diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index 07a61c8587b..a86314e3063 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/tests/unit/proxy/proxy_server/test_routes_queue.py @@ -6,10 +6,13 @@ Pins (PR2): from __future__ import annotations +from collections.abc import Callable +from contextlib import AbstractContextManager from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi.testclient import TestClient from litellm import Router from litellm.proxy import proxy_server @@ -85,6 +88,25 @@ def test_queue_chat_completions_happy(client, auth_as, patched_queue): } +def test_queue_chat_completions_drops_client_sent_stream_usage_flags( + client: TestClient, auth_as: Callable[..., AbstractContextManager[None]], patched_queue: MagicMock +) -> None: + payload: Final = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "priority": 0, + "_litellm_strip_stream_usage": True, + "_litellm_stream_usage_injected": True, + } + with auth_as(): + response: Final = client.post("/queue/chat/completions", json=payload) + assert response.status_code == 200 + forwarded: Final = patched_queue.schedule_acompletion.call_args.kwargs + assert "_litellm_strip_stream_usage" not in forwarded + assert "_litellm_stream_usage_injected" not in forwarded + assert forwarded["messages"] == [{"role": "user", "content": "hi"}] + + def test_queue_chat_completions_no_router_error(client, auth_as, queue_no_router): """Pins ``POST /queue/chat/completions`` (error: no llm_router).""" payload = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 8f27365fcc1..cc61b6f2fcf 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -1430,6 +1430,30 @@ async def test_get_autorouter_presets_fetches_once_per_process(monkeypatch, rese assert calls == ["https://example.test/presets.json"] +@pytest.mark.asyncio +@pytest.mark.parametrize("flag", ["LITELLM_LOCAL_AUTOROUTER_PRESETS", "LITELLM_OFFLINE_MODE"]) +async def test_get_autorouter_presets_local_flags_serve_bundled_without_fetching( + monkeypatch, reset_autorouter_presets_cache, flag +): + from litellm.proxy.public_endpoints.public_endpoints import ( + _load_bundled_autorouter_presets, + get_autorouter_presets, + ) + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + monkeypatch.setenv(flag, "True") + calls = [] + + async def fake_fetch(url): + calls.append(url) + raise AssertionError("no fetch when the bundled catalog is forced") + + presets = await get_autorouter_presets(url="https://example.test/presets.json", fetch=fake_fetch) + + assert presets == _load_bundled_autorouter_presets() + assert calls == [] + + @pytest.mark.asyncio async def test_get_autorouter_presets_single_flight_on_concurrent_cold_start( monkeypatch, reset_autorouter_presets_cache diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index abc1fe88c7a..985aee061bc 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -176,6 +176,105 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert "error" in events[-1] +_COMPACTION_ITEM: Final = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +_HISTORY_BEFORE_COMPACTION: Final = ( + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +) +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "what is the cat's name?"} + + +async def _cursor_stream_failure_prompt_tokens( + monkeypatch: pytest.MonkeyPatch, request_input: list[Mapping[str, str]], fails_after_output: bool +) -> int: + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + class _FailureRowRecorder(CustomLogger): + prompt_tokens: int | None = None + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: object, + traceback_str: str | None = None, + ) -> None: + usage: Final = request_data.get("combined_usage_object") + self.prompt_tokens = usage.prompt_tokens if isinstance(usage, litellm.Usage) else None + + model: Final = "gpt-6-astra" + upstream_url: Final = "https://streaming.example/v1" + created: Final = { + "type": "response.created", + "sequence_number": 0, + "response": { + "id": "resp_upstream", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": model, + "output": [], + }, + } + output_delta: Final = { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_partial", + "output_index": 0, + "content_index": 0, + "delta": "Biscuit", + } + failed: Final = {"type": "error", "error": {"message": "upstream broke", "code": "server_error", "type": None}} + events: Final = (created, output_delta, failed) if fails_after_output else (created, failed) + wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in events) + recorder: Final = _FailureRowRecorder() + router: Final = litellm.Router( + model_list=[ + { + "model_name": model, + "litellm_params": {"model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}, + } + ], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + "/responses").respond( + 200, content=wire, headers={"Content-Type": "text/event-stream"} + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.post( + "/cursor/chat/completions", json={"model": model, "stream": True, "input": request_input} + ) + assert recorder.prompt_tokens is not None, result.text + return recorder.prompt_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fails_after_output", [False, True]) +async def test_a_failed_cursor_stream_counts_its_input_from_the_newest_compaction_item( + monkeypatch: pytest.MonkeyPatch, fails_after_output: bool +) -> None: + """Cursor's route maps to no call type, yet it dispatches a Responses call, which sends + the provider only the newest compaction item onward, so the failure row counts that much.""" + resent: Final = await _cursor_stream_failure_prompt_tokens( + monkeypatch, [*_HISTORY_BEFORE_COMPACTION, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION], fails_after_output + ) + dispatched_after_trim: Final = await _cursor_stream_failure_prompt_tokens( + monkeypatch, [_COMPACTION_ITEM, _TURN_AFTER_COMPACTION], fails_after_output + ) + + assert resent == dispatched_after_trim + + @pytest.mark.asyncio async def test_responses_api_background_polling_rejects_missing_input(): from fastapi import Response as FastAPIResponse @@ -4832,3 +4931,26 @@ def test_responses_routes_document_response_models_in_openapi_schema(): assert "output" in ok_200_properties("/v1/responses/{response_id}", "get") assert "deleted" in ok_200_properties("/v1/responses/{response_id}", "delete") assert "data" in ok_200_properties("/v1/responses/{response_id}/input_items", "get") + + +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_input_token_count_reads_a_compaction_item_as_its_summary(): + from litellm.proxy.response_api_endpoints.endpoints import _responses_input_as_token_count_messages + + assert _responses_input_as_token_count_messages(_COMPACTED_INPUT, None) == ( + _responses_input_as_token_count_messages(_SUMMARY_INPUT, None) + ) diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 979ab499f3b..63e38ce533a 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -1011,10 +1011,9 @@ class TestPollingConditionChecks: # Create mock router mock_router = Mock() - mock_router.model_name_to_deployment_indices = {"gpt-5": [0]} - mock_router.model_list = [ - {"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}} - ] + mock_router.deployments_named.side_effect = { + "gpt-5": ({"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}},) + }.get result = should_use_polling_for_request( background_mode=True, @@ -1033,13 +1032,14 @@ class TestPollingConditionChecks: ) mock_router = Mock() - mock_router.model_name_to_deployment_indices = {"claude-3": [0]} - mock_router.model_list = [ - { - "model_name": "claude-3", - "litellm_params": {"model": "anthropic/claude-3-sonnet"}, - } - ] + mock_router.deployments_named.side_effect = { + "claude-3": ( + { + "model_name": "claude-3", + "litellm_params": {"model": "anthropic/claude-3-sonnet"}, + }, + ) + }.get result = should_use_polling_for_request( background_mode=True, @@ -1183,13 +1183,14 @@ class TestPollingConditionChecks: ) mock_router = Mock() - mock_router.model_name_to_deployment_indices = {"deep-research": [0]} - mock_router.model_list = [ - { - "model_name": "deep-research", - "litellm_params": {"model": "openai/o4-mini-deep-research"}, - } - ] + mock_router.deployments_named.side_effect = { + "deep-research": ( + { + "model_name": "deep-research", + "litellm_params": {"model": "openai/o4-mini-deep-research"}, + }, + ) + }.get # Model alias "deep-research" is in native_background_mode result = should_use_polling_for_request( diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index a1da6c79cd7..71f4dec2345 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from types import MappingProxyType from typing import Final @@ -27,3 +28,59 @@ async def test_large_input_is_still_counted() -> None: assert counts[CL100K_MODEL] == count_input_tokens_for_model(request_body=request_body, model=CL100K_MODEL) assert isinstance(counts, MappingProxyType) + + +_COMPACTION_ITEM: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "And Osaka?"} + + +@pytest.mark.parametrize( + "earlier_item", + [ + pytest.param({"role": "user", "content": "Plan a Tokyo trip. " * 500}, id="user_turn"), + pytest.param({"role": ["system"], "content": "Plan a Tokyo trip."}, id="list_role"), + pytest.param({"role": {"name": "system"}, "content": "Plan a Tokyo trip."}, id="dict_role"), + ], +) +@pytest.mark.parametrize("model", [CL100K_MODEL, "claude-sonnet-5-5"]) +@pytest.mark.asyncio +async def test_history_before_a_compaction_item_is_not_counted(model: str, earlier_item: dict[str, object]) -> None: + compacted_body: Final = {"model": model, "input": [_COMPACTION_ITEM, _TURN_AFTER_COMPACTION]} + full_body: Final = {"model": model, "input": [earlier_item, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION]} + + full_counts: Final = await count_input_tokens( + request_body=full_body, raw_body=json.dumps(full_body).encode(), models=(model,) + ) + compacted_counts: Final = await count_input_tokens( + request_body=compacted_body, raw_body=json.dumps(compacted_body).encode(), models=(model,) + ) + + assert full_counts[model] == compacted_counts[model] + assert count_input_tokens_for_model(request_body=full_body, model=model) == compacted_counts[model] + + +def _compaction_item_carrying(extra_block_fields: dict[str, str]) -> dict[str, str]: + return { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps( + {"type": "compaction", "content": "Turn one planned Tokyo.", **extra_block_fields} + ), + } + + +@pytest.mark.parametrize("model", [CL100K_MODEL, "claude-sonnet-5-5"]) +def test_a_compaction_item_counts_only_its_summary(model: str) -> None: + plain_body: Final = {"model": model, "input": [_compaction_item_carrying({}), _TURN_AFTER_COMPACTION]} + padded_body: Final = { + "model": model, + "input": [_compaction_item_carrying({"opaque_state": "x" * 4000}), _TURN_AFTER_COMPACTION], + } + + assert count_input_tokens_for_model(request_body=padded_body, model=model) == count_input_tokens_for_model( + request_body=plain_body, model=model + ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 23bbe521670..8d53817f51e 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -5,19 +5,21 @@ import hashlib import json import re import sqlite3 +from collections.abc import Mapping from datetime import timezone -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import prisma import pytest -from fastapi import HTTPException +from fastapi import HTTPException, Request from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as ps from litellm.proxy.auth.authorization import OwnedRows from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup, load_permitted_log_team_ids +from litellm.proxy.auth.route_checks import RouteChecks def _default_date_range(): @@ -268,6 +270,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME from litellm.proxy._types import ( + LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, Member, ProxyException, @@ -1894,6 +1898,91 @@ async def test_ui_view_spend_logs_user_filter_intersects_permitted_team_scope(cl app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.parametrize( + ("team_role", "selected_user", "expected_ids"), + ( + ("admin", None, ("member-in-team", "second-member-in-team")), + ("admin", "member", ("member-in-team",)), + ("admin", "caller", ()), + ("admin", "outsider", ()), + ("user", None, None), + (None, "member", None), + ), +) +def test_ui_client_usage_enforces_selected_team_and_member_scope( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + team_role: Literal["admin", "user"] | None, + selected_user: str | None, + expected_ids: tuple[str, ...] | None, +) -> None: + logs: Final = tuple( + { + "request_id": request_id, + "user": user_id, + "team_id": team_id, + "api_key": "shared-key", + "spend": 0.1, + "startTime": "2026-01-02T12:00:00+00:00", + "model": "test-model", + } + for request_id, user_id, team_id in ( + ("member-in-team", "member", "selected-team"), + ("same-member-other-team", "member", "other-team"), + ("second-member-in-team", "second-member", "selected-team"), + ("caller-personal", "caller", None), + ("outsider-request", "outsider", "other-team"), + ) + ) + team: Final = LiteLLM_TeamTable( + team_id="selected-team", + members_with_roles=([Member(user_id="caller", role=team_role)] if team_role else []) + + [Member(user_id="member", role="user"), Member(user_id="second-member", role="user")], + team_member_permissions=[], + ) + + async def team_lookup(where: Mapping[str, object]) -> LiteLLM_TeamTable | None: + return team if where.get("team_id") == team.team_id else None + + def filter_logs(where: Mapping[str, object]) -> tuple[dict[str, str | float | None], ...]: + return tuple( + row + for row in logs + if (where.get("team_id") is None or row["team_id"] == where["team_id"]) + and (where.get("user") is None or row["user"] == where["user"]) + ) + + auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller") + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=LiteLLM_UserTable(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER.value), + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/spend/logs/ui", + request=Request({"type": "http", "method": "GET", "path": "/spend/logs/ui", "query_string": b""}), + valid_token=auth, + request_data={}, + ) + monkeypatch.setattr(ps, "prisma_client", make_ui_spend_logs_mock_prisma(logs, filter_logs, team_lookup)) + monkeypatch.setitem(app.dependency_overrides, ps.user_api_key_auth, lambda: auth) + response: Final = client.get( + "/spend/logs/ui", + params={ + "start_date": "2026-01-01 00:00:00", + "end_date": "2026-01-03 00:00:00", + "team_id": "selected-team", + "page_size": 1000, + "group_by_session": "false", + **({"user_id": selected_user} if selected_user else {}), + }, + headers={"Authorization": "Bearer sk-test"}, + ) + if expected_ids is None: + assert response.status_code == 403 + return + assert response.status_code == 200 + assert tuple(row["request_id"] for row in response.json()["data"]) == expected_ids + assert response.json()["total"] == len(expected_ids) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_pagination(client, monkeypatch): mock_spend_logs = [ @@ -6524,7 +6613,6 @@ def test_ui_view_request_response_reads_from_cold_storage(client, monkeypatch): from litellm.proxy._types import ( LiteLLM_OrganizationMembershipTable, - LiteLLM_UserTable, LiteLLMRoutes, hash_token, ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 55358a3b0de..78f5e1ace0c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -14,8 +14,23 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm import litellm.constants as litellm_constants import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils -from litellm.constants import KUBERNETES_POD_ROUTING_KEY, LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL -from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY +from litellm.constants import ( + KUBERNETES_POD_ROUTING_KEY, + LITELLM_TRUNCATED_PAYLOAD_FIELD, + LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, + LITTELM_CLI_SERVICE_ACCOUNT_NAME, + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + MAX_SPEND_LOG_MODEL_NAME_LENGTH, + REDACTED_BY_LITELM_STRING, + SESSION_ID_GENERATED_METADATA_KEY, + SESSION_ID_OMITTED_METADATA_KEY, + UNKNOWN_MODEL_SPEND_LOG_MODEL, +) +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, +) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -366,7 +381,10 @@ def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(m def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( monkeypatch: pytest.MonkeyPatch, correlation_in_logs: bool, expected_session_id: str ): - header_chain: Final = {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}} + header_chain: Final = { + "litellm_session_id": "trace-abc", + "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}, + } session_id: Final = _chat_session_id( monkeypatch, header_chain, @@ -380,7 +398,11 @@ def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( ("litellm_params", "metadata"), [ pytest.param({"litellm_session_id": "convo\x00x"}, {}, id="body-session"), - pytest.param({"litellm_session_id": "trace-abc"}, {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, id="beside-trace-header"), + pytest.param( + {"litellm_session_id": "trace-abc"}, + {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, + id="beside-trace-header", + ), ], ) def test_caller_session_postgres_cannot_store_keeps_the_trace_id( @@ -394,7 +416,10 @@ def test_caller_session_postgres_cannot_store_keeps_the_trace_id( [ pytest.param({"metadata": {}}, id="no-session"), pytest.param( - {"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + { + "litellm_session_id": "minted-1", + "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, id="minted-by-generate-policy", ), pytest.param({"litellm_session_id": "job-7", "metadata": {"trace_id": "job-7"}}, id="copied-from-trace-id"), @@ -732,6 +757,19 @@ def test_large_request_no_truncation_threshold(): assert "skipped" in truncated_content +@pytest.mark.parametrize( + ("max_string_length_prompt_in_db", "stored_tail"), + [(0, "claude-sonnet-5-5"), (1, "claude-sonnet-5-5"), (-1, "laude-sonnet-5-5"), (-2, "aude-sonnet-5-5")], +) +def test_a_limit_under_two_keeps_mains_tail_after_the_marker(max_string_length_prompt_in_db: int, stored_tail: str): + sanitized: Final = _sanitize_request_body_for_spend_logs_payload( + {"model": "claude-sonnet-5-5"}, max_string_length_prompt_in_db=max_string_length_prompt_in_db + ) + + assert sanitized["model"].startswith("...") + assert sanitized["model"].endswith(f") ...{stored_tail}") + + def test_small_request_no_truncation(): """ Test that small strings are not truncated by MAX_STRING_LENGTH_PROMPT_IN_DB @@ -750,9 +788,7 @@ def test_small_request_no_truncation(): # Verify the content was NOT truncated assert sanitized["messages"][0]["content"] == small_content - assert ( - len(sanitized["messages"][0]["content"]) == MAX_STRING_LENGTH_PROMPT_IN_DB - 100 - ) + assert len(sanitized["messages"][0]["content"]) == MAX_STRING_LENGTH_PROMPT_IN_DB - 100 def test_configurable_string_length_env_var(monkeypatch): @@ -1144,6 +1180,187 @@ def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_ assert parsed["data"][0]["other_field"] == "value" +def test_get_response_for_spend_logs_payload_keeps_a_long_compaction_summary_whole(monkeypatch: pytest.MonkeyPatch): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + summary = f"{filler}The cat is named Biscuit. {filler}" + answer = f"{filler}OK{filler}" + payload = cast( + StandardLoggingPayload, + { + "response": { + "choices": [ + { + "message": { + "role": "assistant", + "content": answer, + "provider_specific_fields": { + "compaction_blocks": [{"type": "compaction", "content": summary}] + }, + } + } + ] + } + }, + ) + + message = json.loads(_get_response_for_spend_logs_payload(payload))["choices"][0]["message"] + assert message["provider_specific_fields"]["compaction_blocks"] == [{"type": "compaction", "content": summary}] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message["content"] + + +def test_stored_request_keeps_only_the_summary_of_a_long_claude_compaction_input_item(monkeypatch: pytest.MonkeyPatch): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + summary = f"{filler}The cat is named Biscuit. {filler}" + claude_item = json.dumps({"type": "compaction", "content": summary, "padding": filler * 10}) + opaque_item = "gAAAA" + "x" * (MAX_STRING_LENGTH_PROMPT_IN_DB * 2) + request_body = { + "input": [ + {"type": "compaction", "id": "cmp_1", "encrypted_content": claude_item}, + {"type": "compaction", "id": "cmp_2", "encrypted_content": opaque_item}, + {"role": "assistant", "content": [{"type": "compaction", "content": f"{filler}Biscuit{filler}"}]}, + {"role": "user", "content": f"{filler}What is my cat called?{filler}"}, + ], + "metadata": {"note": {"type": "compaction", "encrypted_content": claude_item}}, + } + + stored = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params={"proxy_server_request": {"body": request_body}} + ) + ) + + claude_stored, opaque_stored, messages_block_stored, message_stored = stored["input"] + assert {**claude_stored, "encrypted_content": json.loads(claude_stored["encrypted_content"])} == { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": {"type": "compaction", "content": summary}, + } + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in opaque_stored["encrypted_content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in messages_block_stored["content"][0]["content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message_stored["content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored["metadata"]["note"]["encrypted_content"] + + +def _compaction_summary_past_the_summary_cap() -> str: + from litellm.constants import MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + + filler: Final = "detail " * (MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB // 7) + return f"The cat is named Biscuit. {filler}The dog is named Rex." + + +def _assert_bounded_like_a_db_string(stored: str, summary: str) -> None: + from litellm.constants import MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored + assert len(stored) < MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + 500 < len(summary) + 500 + assert stored.startswith("The cat is named Biscuit.") + assert stored.endswith("The dog is named Rex.") + + +@pytest.mark.parametrize("operator_cap_above_the_summary", [False, True]) +def test_spend_row_bounds_a_compaction_summary_at_the_larger_of_the_summary_cap_and_the_operator_cap( + monkeypatch: pytest.MonkeyPatch, operator_cap_above_the_summary: bool +): + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + summary: Final = _compaction_summary_past_the_summary_cap() + if operator_cap_above_the_summary: + monkeypatch.setenv("MAX_STRING_LENGTH_PROMPT_IN_DB", str(len(summary) * 2)) + payload: Final = cast( + StandardLoggingPayload, + { + "response": { + "choices": [ + { + "message": { + "role": "assistant", + "content": "OK", + "provider_specific_fields": { + "compaction_blocks": [{"type": "compaction", "content": summary}] + }, + } + } + ] + } + }, + ) + request_body: Final = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary}), + } + ] + } + + response_message: Final = json.loads(_get_response_for_spend_logs_payload(payload))["choices"][0]["message"] + stored_item: Final = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params={"proxy_server_request": {"body": request_body}} + ) + )["input"][0] + + stored_response_block: Final = response_message["provider_specific_fields"]["compaction_blocks"][0] + stored_request_block: Final = json.loads(stored_item["encrypted_content"]) + assert stored_request_block["type"] == stored_response_block["type"] == "compaction" + if operator_cap_above_the_summary: + assert stored_response_block["content"] == stored_request_block["content"] == summary + return + _assert_bounded_like_a_db_string(stored_response_block["content"], summary) + _assert_bounded_like_a_db_string(stored_request_block["content"], summary) + + +@pytest.mark.parametrize("summary_past_the_summary_cap", [False, True]) +def test_responses_shaped_spend_row_keeps_a_compaction_output_item_summary_bounded_like_a_chat_block( + monkeypatch: pytest.MonkeyPatch, summary_past_the_summary_cap: bool +): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + from litellm.responses.litellm_completion_transformation.compaction_items import compaction_summary + from litellm.responses.utils import ResponsesAPIRequestUtils + + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + filler: Final = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + summary: Final = ( + _compaction_summary_past_the_summary_cap() + if summary_past_the_summary_cap + else f"{filler}The cat is named Biscuit. {filler}" + ) + payload: Final = cast( + StandardLoggingPayload, + { + "response": { + "object": "response", + "output": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + json.dumps({"type": "compaction", "content": summary}), "deployment-1" + ), + }, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]}, + ], + } + }, + ) + + stored_item: Final = json.loads(_get_response_for_spend_logs_payload(payload))["output"][0] + model_id, _ = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(stored_item["encrypted_content"]) + stored_summary: Final = compaction_summary(stored_item) + assert model_id == "deployment-1" + assert stored_summary is not None + if summary_past_the_summary_cap: + _assert_bounded_like_a_db_string(stored_summary, summary) + return + assert stored_summary == summary + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from response.""" diff --git a/tests/unit/proxy/telemetry/test_endpoints.py b/tests/unit/proxy/telemetry/test_endpoints.py index 3b5f66da0ed..c962c8df6a5 100644 --- a/tests/unit/proxy/telemetry/test_endpoints.py +++ b/tests/unit/proxy/telemetry/test_endpoints.py @@ -1,13 +1,12 @@ -from dataclasses import dataclass, field from collections.abc import Mapping +from dataclasses import dataclass, field from typing import Final -from typing_extensions import LiteralString - import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from httpx import ASGITransport, AsyncClient, Response +from typing_extensions import LiteralString from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth diff --git a/tests/unit/proxy/telemetry/test_settings.py b/tests/unit/proxy/telemetry/test_settings.py index 0e9c180da4f..b88e6811113 100644 --- a/tests/unit/proxy/telemetry/test_settings.py +++ b/tests/unit/proxy/telemetry/test_settings.py @@ -1,7 +1,6 @@ from typing import Final import pytest - from pydantic import ValidationError from litellm.proxy.telemetry.settings import EnvPolicy, TelemetrySettings, describe_errors, env_policy, load_settings diff --git a/tests/unit/proxy/test__lazy_features.py b/tests/unit/proxy/test__lazy_features.py index adc4de8c8fa..0749fc12ea9 100644 --- a/tests/unit/proxy/test__lazy_features.py +++ b/tests/unit/proxy/test__lazy_features.py @@ -93,6 +93,20 @@ def test_retired_roi_calculator_cannot_be_loaded_or_called( assert response.status_code == 404, f"{method} {path}: {response.text}" +@pytest.mark.parametrize("disable_lazy_routes", ("false", "true")) +@pytest.mark.parametrize("path", ("/lazy/warm/usage_ai", "/usage/ai/chat")) +def test_retired_usage_chat_cannot_be_loaded_or_called( + monkeypatch: pytest.MonkeyPatch, disable_lazy_routes: str, path: str +) -> None: + monkeypatch.setenv(FLAG, disable_lazy_routes) + app: Final = FastAPI() + attach_lazy_features(app) + + with TestClient(app) as client: + response: Final = client.post(path) + assert response.status_code == 404, f"{path}: {response.text}" + + @pytest.mark.parametrize("value", ("1", "true", "TRUE", "yes", "on")) def test_flag_registers_every_feature_at_startup(monkeypatch: pytest.MonkeyPatch, value: str) -> None: monkeypatch.setenv(FLAG, value) diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 68bfe4b34e5..76ca49362b4 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -2,23 +2,35 @@ import json from typing import Final import pytest -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from litellm.proxy._types import ( ROLES_WITHIN_ORG, + BudgetNewRequest, ChangePasswordRequest, GenerateKeyRequest, + GenerateKeyResponse, KeyRequest, LiteLLM_AuditLogs, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, + NewCustomerRequest, NewMCPServerRequest, + NewTeamRequest, NewUserRequest, + NewUserResponse, OrganizationMemberUpdateRequest, + RegenerateKeyRequest, ResetSpendRequest, UpdateKeyRequest, UpdateMCPServerRequest, + UpdateTeamRequest, UpdateUserRequest, + UpdateUserRequestNoUserIDorEmail, UserAPIKeyAuth, ) @@ -213,6 +225,75 @@ def test_a_max_budget_that_was_actually_sent_is_kept(sent): assert GenerateKeyRequest(max_budget=sent).max_budget == sent +RATE_LIMIT_REQUESTS: Final = ( + (GenerateKeyRequest, {}), + (RegenerateKeyRequest, {}), + (UpdateKeyRequest, {"key": "sk-test"}), + (NewUserRequest, {}), + (UpdateUserRequestNoUserIDorEmail, {}), + (UpdateUserRequest, {"user_id": "user-test"}), + (NewTeamRequest, {}), + (UpdateTeamRequest, {"team_id": "team-test"}), + (BudgetNewRequest, {}), + (NewCustomerRequest, {"user_id": "customer-test"}), +) + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +@pytest.mark.parametrize("field", ("rpm_limit", "tpm_limit")) +@pytest.mark.parametrize("value", (-1, -100)) +def test_management_requests_reject_negative_rate_limits( + request_model: type[BaseModel], identifiers: dict[str, str], field: str, value: int +) -> None: + with pytest.raises(ValidationError) as error: + request_model.model_validate({**identifiers, field: value}) + assert error.value.errors()[0]["loc"] == (field,) + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +@pytest.mark.parametrize("value", (None, 0, 100)) +def test_management_requests_preserve_explicit_rate_limits( + request_model: type[BaseModel], identifiers: dict[str, str], value: int | None +) -> None: + limits: Final = {"rpm_limit": value, "tpm_limit": value} + request: Final = request_model.model_validate({**identifiers, **limits}) + assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == limits + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +def test_management_requests_leave_omitted_rate_limits_unset( + request_model: type[BaseModel], identifiers: dict[str, str] +) -> None: + request: Final = request_model.model_validate({**identifiers}) + assert request.model_dump(include={"rpm_limit", "tpm_limit"}) == {"rpm_limit": None, "tpm_limit": None} + assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == {} + + +@pytest.mark.parametrize(("response_model", "identifiers"), ( + (GenerateKeyResponse, {"key": "sk-test"}), + (NewUserResponse, {"key": "sk-test"}), + (LiteLLM_TeamTable, {"team_id": "team-test"}), + (LiteLLM_UserTable, {"user_id": "user-test"}), + (LiteLLM_BudgetTable, {"budget_id": "budget-test"}), +)) +def test_existing_negative_rate_limits_remain_readable( + response_model: type[BaseModel], identifiers: dict[str, str] +) -> None: + limits: Final = {"rpm_limit": -1, "tpm_limit": -100} + response: Final = response_model.model_validate({**identifiers, **limits}) + assert response.model_dump(include={"rpm_limit", "tpm_limit"}) == limits + + +def test_existing_customer_negative_rate_limits_remain_readable() -> None: + customer: Final = LiteLLM_EndUserTable.model_validate({ + "user_id": "customer-test", "blocked": False, + "litellm_budget_table": {"rpm_limit": -1, "tpm_limit": -100}, + }) + assert customer.litellm_budget_table is not None + assert customer.litellm_budget_table.rpm_limit == -1 + assert customer.litellm_budget_table.tpm_limit == -100 + + USER_IDENTIFIER_CASES = ( ({"user_id": "user-1"}, True), ({"user_email": "user@example.com"}, True), diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index d742956bb59..82dd3479321 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4066,6 +4066,72 @@ class TestHasAttributeErrorInChain: assert _has_attribute_error_in_chain(exc_a) is False +_COMPACTION_ITEM: Final = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +_HISTORY_BEFORE_COMPACTION: Final = ( + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +) +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "what is the cat's name?"} + + +class TestHandleLLMApiExceptionPricesWhatTheServerDispatched: + async def _failed_prompt_tokens( + self, + route_type: Literal["acompletion", "aresponses"], + logged_call_type: str, + request_input: list[dict[str, str]], + ) -> int: + processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4o", "input": request_input}) + cursor_caller: Final = ProxyUserAPIKeyAuth(request_route="/cursor/chat/completions") + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alert_types = [] + _, logging_obj = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=cursor_caller, + proxy_logging_obj=proxy_logging, + proxy_config=MagicMock(spec=ProxyConfig), + route_type=route_type, + ) + logging_obj.model_call_details.update( + first_api_call_start_time=datetime.datetime.now(), + messages=request_input, + call_type=logged_call_type, + ) + with pytest.raises(ProxyException): + await processing_obj.handle_llm_api_exception( + e=Exception("boom"), + user_api_key_dict=cursor_caller, + proxy_logging_obj=proxy_logging, + ) + return processing_obj.data["combined_usage_object"].prompt_tokens + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "route_type, logged_call_type, trims_at_the_compaction", + [("aresponses", "aresponses", True), ("acompletion", "responses", False)], + ) + async def test_a_failed_call_counts_from_the_newest_compaction_only_when_dispatched_as_responses( + self, + route_type: Literal["acompletion", "aresponses"], + logged_call_type: str, + trims_at_the_compaction: bool, + ) -> None: + async def prompt_tokens(request_input: list[dict[str, str]]) -> int: + return await self._failed_prompt_tokens(route_type, logged_call_type, request_input) + + resent: Final = await prompt_tokens([*_HISTORY_BEFORE_COMPACTION, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION]) + dispatched_after_trim: Final = await prompt_tokens([_COMPACTION_ITEM, _TURN_AFTER_COMPACTION]) + + assert (resent == dispatched_after_trim) is trims_at_the_compaction + assert resent >= dispatched_after_trim + + @pytest.mark.asyncio class TestHandleLLMApiExceptionDictDetail: """ @@ -9181,6 +9247,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_usage": True} assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_default_preserves_other_client_stream_options_keys(self): data = {"stream": True, "stream_options": {"include_obfuscation": True}} @@ -9189,6 +9256,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_obfuscation": True, "include_usage": True} assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_client_requested_usage_is_left_untouched_and_not_stripped(self): data = {"stream": True, "stream_options": {"include_usage": True}} @@ -9205,6 +9273,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"]["include_usage"] is True assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_explicit_false_flag_disables_injection_entirely(self): data = {"stream": True} @@ -9230,6 +9299,36 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_usage": True} assert "_litellm_strip_stream_usage" not in data + @pytest.mark.parametrize("client_stream_options", [None, {"include_obfuscation": True}]) + def test_flag_true_marks_usage_injected_for_clients_that_did_not_ask( + self, client_stream_options: dict[str, bool] | None + ): + data: Final = { + "stream": True, + **({} if client_stream_options is None else {"stream_options": client_stream_options}), + } + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"]["include_usage"] is True + assert data["_litellm_stream_usage_injected"] is True + + def test_flag_true_leaves_usage_unmarked_when_client_asked_and_neutralizes_client_marker(self): + data: Final = {"stream": True, "stream_options": {"include_usage": True}, "_litellm_stream_usage_injected": True} + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"] == {"include_usage": True} + assert data["_litellm_stream_usage_injected"] is False + def test_flag_true_respects_client_explicit_include_usage_false(self): data = {"stream": True, "stream_options": {"include_usage": False}} @@ -10784,6 +10883,7 @@ class TestDetachedStreamFailureHook: "user_api_key_dict": user_api_key_dict, "original_exception": failure, "request_data": processing_obj.data, + "dispatched_call_type": "anthropic_messages", } ] @@ -11927,3 +12027,24 @@ class TestResponseCostFromLoggingObjForCacheHits: ) assert cost == "" + + +def test_maybe_get_model_id_reads_the_chat_metadata_the_router_updated_in_place(): + processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"metadata": {"model_info": {"id": "declared-text"}}}) + + assert processing_obj.maybe_get_model_id(None) == "declared-text" + + +def test_maybe_get_model_id_prefers_litellm_metadata_over_caller_metadata_without_a_deployment(): + processing_obj: Final = ProxyBaseLLMRequestProcessing( + data={ + "metadata": {"trace_id": "t-1"}, + "litellm_metadata": {"model_info": {"id": "vision"}}, + } + ) + + assert processing_obj.maybe_get_model_id(None) == "vision" + + +def test_maybe_get_model_id_returns_none_before_the_router_picked_a_deployment(): + assert ProxyBaseLLMRequestProcessing(data={"metadata": {}}).maybe_get_model_id(None) is None diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c2a5b01fe4e..ed7817bd660 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1,4 +1,5 @@ import asyncio +import base64 import copy import json import os @@ -106,6 +107,13 @@ class TestGetMetadataVariableName: request = self._make_request("/v1/files") assert get_metadata_variable_name(request) == "litellm_metadata" + @pytest.mark.parametrize( + "path", ["/v1beta/interactions", "/interactions", "/v1beta/interactions/v1_abc123", "/interactions/v1_abc123/cancel"] + ) + def test_returns_litellm_metadata_for_interactions_routes(self, path): + request = self._make_request(path) + assert get_metadata_variable_name(request) == "litellm_metadata" + def test_returns_metadata_for_chat_completions(self): request = self._make_request("/chat/completions") assert get_metadata_variable_name(request) == "metadata" @@ -708,6 +716,46 @@ async def test_add_litellm_data_to_request_parses_json_object_string_litellm_met assert updated["litellm_metadata"]["cost_centre"] == "research" +@pytest.mark.asyncio +async def test_interactions_spend_logs_metadata_survives_the_router_metadata_bucket(): + """GH#47841: on /v1beta/interactions the router plants its deployment fields in + `litellm_metadata`. The proxy must put the request's metadata there too: otherwise the + spend log reads only the router's bucket and `spend_logs_metadata` from the + `x-litellm-spend-logs-metadata` header never reaches LiteLLM_SpendLogs.""" + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1beta/interactions" + request_mock.url.path = "/v1beta/interactions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = { + "Content-Type": "application/json", + "x-litellm-spend-logs-metadata": '{"cost_centre": "research"}', + } + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + request_mock.state.parent_otel_span = None + + data = await add_litellm_data_to_request( + data={"model": "gemini-3.1-flash-lite", "input": "Reply OK.", "store": True}, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert data["litellm_metadata"]["spend_logs_metadata"] == {"cost_centre": "research"} + + # The router then adds its deployment fields to the same bucket, and the spend log + # resolves that bucket. + data["litellm_metadata"].update({"model_group": "gemini-3.1-flash-lite", "deployment": "gemini/gemini-3.1-flash-lite"}) + resolved = get_litellm_metadata_from_kwargs( + {"litellm_params": {"metadata": data.get("metadata") or {}, "litellm_metadata": data["litellm_metadata"]}} + ) + assert resolved["spend_logs_metadata"] == {"cost_centre": "research"} + assert resolved["model_group"] == "gemini-3.1-flash-lite" + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_proxy_server_request_body_is_post_strip(): """Regression: proxy_server_request['body'] used to be snapshotted before @@ -936,6 +984,26 @@ def test_body_snapshot_excludes_team_callback_credentials() -> None: }, proxy_request +def test_body_snapshot_keeps_the_stream_usage_injection_marker_out_of_the_stored_request() -> None: + from litellm.proxy.common_request_processing import stream_usage_tracking_updates + from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot + from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD + + request: Final = {"model": "gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True} + proxy_request: Final = {"body": dict(request)} + data: Final = { + **request, + **stream_usage_tracking_updates(request, {}, "acompletion", lambda: True), + "proxy_server_request": proxy_request, + } + + refresh_proxy_server_request_body_snapshot(data) + + assert data[STREAM_USAGE_INJECTED_FIELD] is True + assert proxy_request["body"]["stream_options"] == {"include_usage": True} + assert STREAM_USAGE_INJECTED_FIELD not in proxy_request["body"] + + @pytest.mark.asyncio @pytest.mark.parametrize("pre_call_ran", [False, True]) async def test_post_guardrail_snapshot_preserves_logging_only_masking_in_spend_logs( @@ -1209,6 +1277,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "litellm_metadata": copy.deepcopy(malicious_metadata), "weights": {"gpt-3.5-turbo": {"forged-deployment-id": 100}}, "_router_weights": {"gpt-3.5-turbo": {"forged-deployment-id": 100}}, + "_encrypted_content_affinity_pinned": True, } updated = await add_litellm_data_to_request( @@ -1231,6 +1300,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): assert "_router_weights" not in updated assert "weights" not in updated["proxy_server_request"]["body"] assert "_router_weights" not in updated["proxy_server_request"]["body"] + assert "_encrypted_content_affinity_pinned" not in updated stripped_keys = { "disable_global_guardrails", @@ -8042,6 +8112,67 @@ def test_client_anthropic_api_headers_stay_off_openai_compatible_providers(): assert forwarded == {} +def test_client_anthropic_api_headers_stay_off_azure_ai_outside_native_messages(): + forwarded = _headers_forwarded_to({"anthropic-beta": "thinking-binding-controls-2026-08-01"}, "azure_ai") + + assert forwarded == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_client_thinking_binding_beta_reaches_azure_ai_foundry_claude(): + import httpx + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + data: Final[dict] = {"model": "claude-opus-4-8"} # mutable-ok: the function under test writes into it + add_provider_specific_headers_to_request( + data=data, + headers={ + "content-type": "application/json", + "anthropic-beta": "thinking-binding-controls-2026-08-01", + "anthropic-version": "2026-08-01", + }, + ) + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake Foundry endpoint records what the proxy sent + + def foundry_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response( + status_code=200, + json={ + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + thinking={"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "drop_block"}}, + model="azure_ai/claude-opus-4-8", + api_key="foundry-key", + api_base="https://my-resource.services.ai.azure.com/anthropic", + client=upstream, + provider_specific_header=data["provider_specific_header"], + ) + finally: + await upstream.close() + + assert received[0].headers["anthropic-beta"] == "thinking-binding-controls-2026-08-01" + assert received[0].headers["anthropic-version"] == "2026-08-01" + assert json.loads(received[0].content)["thinking"]["block_binding"] == {"prefix_mismatch_behavior": "drop_block"} + + TYPESAFE_DANGER_ZONE_HEADERS = { "TypeSafe-Danger-Zone-Images": "true", "typesafe-danger-zone-long-context": "true", @@ -9078,6 +9209,97 @@ def test_default_team_settings_bool_turn_off_message_logging_redacts(): ) +def _fake_chatgpt_oauth_token(account_id: str) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode( + json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}, "exp": 4102444800}).encode() + ) + .rstrip(b"=") + .decode() + ) + return f"{header}.{payload}.fake-signature" + + +CHATGPT_OAUTH_TOKEN = f"Bearer {_fake_chatgpt_oauth_token('acct-alice')}" +CHATGPT_ACCOUNT_ID_HEADER_NAME = "ChatGPT-Account-Id" +PLATFORM_API_KEY_BEARER = "Bearer sk-proj-fake-platform-key-0123456789" +CHATGPT_LEAK_TARGET_PROVIDERS = ["openai", "azure", "anthropic", "bedrock", "vertex_ai"] + + +def _codex_client_headers( + authorization_header_name: str = "authorization", account_id: str | None = "acct-alice" +) -> dict: + return { + "content-type": "application/json", + "originator": "codex_cli_rs", + authorization_header_name: CHATGPT_OAUTH_TOKEN, + **({CHATGPT_ACCOUNT_ID_HEADER_NAME: account_id} if account_id is not None else {}), + } + + +def test_clean_headers_keeps_chatgpt_oauth_bearer_when_gateway_key_is_separate(): + result = clean_headers( + headers=Headers({**_codex_client_headers(), "x-litellm-api-key": "sk-proxy-virtual-key"}), + litellm_key_header_name="x-litellm-api-key", + forward_llm_provider_auth_headers=False, + authenticated_with_header="x-litellm-api-key", + ) + + assert result["authorization"] == CHATGPT_OAUTH_TOKEN + assert "x-litellm-api-key" not in result + + +def test_clean_headers_drops_bearer_that_authenticated_to_the_gateway(): + result = clean_headers( + headers=Headers({"authorization": CHATGPT_OAUTH_TOKEN, "content-type": "application/json"}), + litellm_key_header_name=None, + forward_llm_provider_auth_headers=False, + authenticated_with_header="authorization", + ) + + assert _authorization_values(result) == [] + + +@pytest.mark.parametrize("bearer", [PLATFORM_API_KEY_BEARER, "Bearer not.a-jwt", "Bearer a.bm90LWpzb24.c"]) +def test_clean_headers_drops_non_chatgpt_bearers(bearer): + result = clean_headers( + headers=Headers({"authorization": bearer, "x-litellm-api-key": "sk-proxy-virtual-key"}), + litellm_key_header_name="x-litellm-api-key", + forward_llm_provider_auth_headers=False, + authenticated_with_header="x-litellm-api-key", + ) + + assert _authorization_values(result) == [] + + +@pytest.mark.parametrize("authorization_header_name", AUTHORIZATION_HEADER_CASINGS) +def test_chatgpt_oauth_credential_and_account_reach_chatgpt_provider(authorization_header_name): + forwarded = _headers_forwarded_to(_codex_client_headers(authorization_header_name), "chatgpt") + + assert _authorization_values(forwarded) == [CHATGPT_OAUTH_TOKEN] + assert forwarded[CHATGPT_ACCOUNT_ID_HEADER_NAME] == "acct-alice" + + +@pytest.mark.parametrize("custom_llm_provider", CHATGPT_LEAK_TARGET_PROVIDERS) +def test_chatgpt_oauth_credential_is_never_forwarded_to_other_providers(custom_llm_provider): + forwarded = _headers_forwarded_to(_codex_client_headers(), custom_llm_provider) + + assert _authorization_values(forwarded) == [] + assert CHATGPT_OAUTH_TOKEN not in forwarded.values() + assert CHATGPT_ACCOUNT_ID_HEADER_NAME not in forwarded + + +def test_chatgpt_account_header_is_not_forwarded_without_chatgpt_oauth_bearer(): + headers = {**_codex_client_headers(), "authorization": PLATFORM_API_KEY_BEARER} + + assert _headers_forwarded_to(headers, "chatgpt") == {} + + +def test_anthropic_oauth_bearer_does_not_reach_chatgpt_provider(): + assert _authorization_values(_headers_forwarded_to(_client_headers(), "chatgpt")) == [] + + def test_add_user_api_key_auth_to_request_metadata_attributes_a_cli_session_to_its_alias(): data = {"model": "gpt-5.4-nano", "messages": [{"role": "user", "content": "hi"}], "litellm_metadata": {}} session = UserAPIKeyAuth( diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index f8ea1e5d075..a26a64c6fb0 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -3813,3 +3813,47 @@ def test_live_session_preserves_pretty_printed_start_frame( assert captured.route_type == "_arealtime" assert routed_data["model"] == "gpt-live-1" assert routed_data["live_session_start"] == expected_start + + +@pytest.mark.parametrize("scope", [None, "expand"]) +@pytest.mark.parametrize( + "general_settings, wildcard_query, wildcard_listed", + [ + ({}, None, False), + ({"return_wildcard_routes": False}, None, False), + ({"return_wildcard_routes": True}, None, True), + ({"return_wildcard_routes": "true"}, None, False), + ({"return_wildcard_routes": True}, False, False), + ({"return_wildcard_routes": False}, True, True), + ({}, True, True), + ], +) +@pytest.mark.asyncio +async def test_model_list_return_wildcard_routes_defaults_to_general_setting( + monkeypatch: pytest.MonkeyPatch, + general_settings: dict[str, object], + wildcard_query: bool | None, + wildcard_listed: bool, + scope: str | None, +) -> None: + router: Final = litellm.Router( + model_list=[ + {"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": "sk-fake"}}, + {"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": "sk-fake"}}, + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_model_list", router.model_list) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", None) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_model", None) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", general_settings) + + response: Final = await litellm.proxy.proxy_server.model_list( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=LitellmUserRoles.PROXY_ADMIN), + return_wildcard_routes=wildcard_query, + scope=scope, + ) + listed: Final = [model["id"] for model in response["data"]] + + assert "gpt-6.1-sol" in listed, listed + assert ("openai/*" in listed) is wildcard_listed, listed 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 14196d4a1ac..a859ce96c70 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -15207,21 +15207,23 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun assert_loop_stayed_free(took, lags) -async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch): - from litellm.rust_bridge._native import Tokenizer +async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None: + from queue import SimpleQueue + from threading import Thread, current_thread + + from tokenizers import Tokenizer from litellm import Router - from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags - claude_tokenizer: Final = litellm.utils.select_tokenizer("claude-fable-5")["tokenizer"] + claude_tokenizer: Final = Tokenizer.from_str(litellm.utils.claude_json_str) + loop_thread: Final = current_thread() + loader_threads: Final[SimpleQueue[Thread]] = SimpleQueue() - class SlowHubTokenizer: - @staticmethod - def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: - time.sleep(0.3) - return claude_tokenizer + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + loader_threads.put(current_thread()) + return claude_tokenizer - monkeypatch.setattr("litellm.rust_bridge.tokenizer.from_pretrained", SlowHubTokenizer.from_pretrained) + monkeypatch.setattr("tokenizers.Tokenizer", types.SimpleNamespace(from_pretrained=from_pretrained)) monkeypatch.setattr( "litellm.proxy.proxy_server.llm_router", Router( @@ -15237,15 +15239,16 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp ), ) - response, took, lags = await timed_with_loop_lags( - lambda: proxy_server_module.token_counter( + litellm.utils._select_custom_tokenizer_helper.cache_clear() + try: + response: Final = await proxy_server_module.token_counter( TokenCountRequest(model="self-hosted", prompt="count me off the loop") ) - ) - - assert response.tokenizer_type == "huggingface_tokenizer" - assert response.total_tokens > 0 - assert_loop_stayed_free(took, lags) + assert loader_threads.get_nowait() is not loop_thread + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens > 0 + finally: + litellm.utils._select_custom_tokenizer_helper.cache_clear() async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch): diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 7b18cb3350c..b5940e4c212 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1,5 +1,9 @@ +import asyncio import datetime as real_datetime +import json import smtplib +import traceback +from collections.abc import Mapping from typing import Final import httpx @@ -11,8 +15,15 @@ from litellm.exceptions import InternalServerError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.bug_report import ISSUE_URL_BASE +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth -from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy +from litellm.proxy.utils import ( + PrismaClient, + ProxyLogging, + _failure_fields_to_lift, + handle_exception_on_proxy, +) +from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage from litellm.types.guardrails import GuardrailEventHooks @@ -735,7 +746,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: post-dispatch failures keep zero spend. """ - async def _run(self, request_data, original_exception: Exception | None = None): + async def _run( + self, request_data, original_exception: Exception | None = None, request_route=None, dispatched_call_type=None + ): from unittest.mock import AsyncMock, patch from litellm.proxy._types import UserAPIKeyAuth @@ -746,7 +759,8 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: await proxy_logging_obj.post_call_failure_hook( request_data=request_data, original_exception=original_exception or Exception("boom"), - user_api_key_dict=UserAPIKeyAuth(), + user_api_key_dict=UserAPIKeyAuth(request_route=request_route), + dispatched_call_type=dispatched_call_type, ) def _timeout(self, model: str, cause: Exception) -> Exception: @@ -763,6 +777,38 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: logging_obj.model_call_details = model_call_details return logging_obj + @pytest.mark.asyncio + @pytest.mark.parametrize( + "request_route, dispatched_call_type", + [("/v1/responses", None), ("/cursor/chat/completions", "aresponses")], + ) + async def test_a_failed_responses_request_counts_its_input_from_the_newest_compaction_item( + self, request_route, dispatched_call_type + ): + from datetime import datetime + + async def prompt_tokens(request_input): + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": FAILURE_USAGE_MODEL, + "messages": request_input, + "call_type": "aresponses", + } + ), + "metadata": {}, + } + await self._run(request_data, request_route=request_route, dispatched_call_type=dispatched_call_type) + return request_data["combined_usage_object"].prompt_tokens + + resent = await prompt_tokens( + [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + + assert resent == await prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + assert resent < await prompt_tokens([COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, TURN_AFTER_COMPACTION]) + @pytest.mark.asyncio async def test_timed_out_request_estimates_input_tokens_priced_at_input_rate(self): from datetime import datetime @@ -2154,29 +2200,29 @@ def test_a_system_prompt_reads_the_same_whatever_shape_it_arrived_in(system_inpu def test_a_system_prompt_is_counted_on_top_of_the_request(): - from litellm.proxy.utils import _count_request_input_tokens + from litellm.proxy.utils import count_request_input_tokens - without = _count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", None) - with_system = _count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", "be brief") + without = count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", None) + with_system = count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", "be brief") assert without > 0 assert with_system > without def test_a_request_with_nothing_in_it_counts_zero(): - from litellm.proxy.utils import _count_request_input_tokens + from litellm.proxy.utils import count_request_input_tokens - assert _count_request_input_tokens(FAILURE_USAGE_MODEL, [], None) == 0 - assert _count_request_input_tokens(FAILURE_USAGE_MODEL, None, None) == 0 + assert count_request_input_tokens(FAILURE_USAGE_MODEL, [], None) == 0 + assert count_request_input_tokens(FAILURE_USAGE_MODEL, None, None) == 0 def test_a_failed_dispatch_is_estimated_as_input_only(): - from litellm.proxy.utils import _count_request_input_tokens, _estimate_dispatched_failure_usage + from litellm.proxy.utils import count_request_input_tokens, _estimate_dispatched_failure_usage usage = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage is not None - assert usage.prompt_tokens == _count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) + assert usage.prompt_tokens == count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage.completion_tokens == 0 assert usage.total_tokens == usage.prompt_tokens @@ -2440,7 +2486,7 @@ def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_stre lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) - estimate = _estimated_prompt_usage(logging_obj.model_call_details, {}) + estimate = _estimated_prompt_usage(logging_obj.model_call_details, {}, None) assert lifted["combined_usage_object"] is recovered assert [(a.usage.prompt_tokens, a.usage.completion_tokens) for a in lifted["billed_failed_attempts"]] == [ (estimate.prompt_tokens, 0) @@ -2472,10 +2518,681 @@ def test_attempts_already_summed_into_the_final_row_are_not_lifted_again(): lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) - assert lifted["combined_usage_object"].prompt_tokens == 24 + current_only_logging_obj: Final = _dispatched_failure_logging_obj( + FAILURE_USAGE_MODEL, + get_litellm_params( + metadata={ + "attempted_retries": 1, + "model_info": {"id": "dep-lit145-final"}, + "previous_models": ( + {"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0}, + ), + } + ), + ) + current_only: Final = _failure_fields_to_lift( + {"litellm_logging_obj": current_only_logging_obj}, provider_billed=True + ) + current_only_usage: Final = current_only["combined_usage_object"] + + assert current_only_usage.prompt_tokens > 0 + assert lifted["combined_usage_object"].prompt_tokens == 24 + current_only_usage.prompt_tokens assert "billed_failed_attempts" not in lifted +def test_carry_superseded_attempt_usage_without_stashed_usage_preserves_failure_estimate() -> None: + request_start: Final = real_datetime.datetime(2026, 1, 1) + logging_obj: Final = Logging( + model=FAILURE_USAGE_MODEL, + messages=ONE_USER_MESSAGE, + stream=False, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-no-stashed-usage", + function_id="fn-lit145-no-stashed-usage", + ) + logging_obj.pre_call(input=ONE_USER_MESSAGE, api_key="sk-test", model=FAILURE_USAGE_MODEL) + logging_obj.model_call_details["call_type"] = "acompletion" + logging_obj.model_call_details["messages"] = ONE_USER_MESSAGE + + try: + try: + raise httpx.ReadTimeout("no bytes in time") + except httpx.ReadTimeout as cause: + raise litellm.Timeout("Request timed out.", model=FAILURE_USAGE_MODEL, llm_provider="openai") from cause + except litellm.Timeout as timeout: + logging_obj.failure_handler( + timeout, + traceback.format_exc(), + request_start, + request_start + real_datetime.timedelta(seconds=1), + ) + + assert "combined_usage_object" not in logging_obj.model_call_details + request_data: Final = {"litellm_logging_obj": logging_obj} + before: Final = _failure_fields_to_lift(request_data, provider_billed=True) + + carry_superseded_attempt_usage(request_data) + + after: Final = _failure_fields_to_lift(request_data, provider_billed=True) + before_usage: Final = before.get("combined_usage_object") + after_usage: Final = after.get("combined_usage_object") + assert before_usage is not None + assert after_usage is not None + assert after_usage == before_usage + assert before["response_cost"] > 0 + assert after["response_cost"] == before["response_cost"] + + +@pytest.mark.asyncio +async def test_dropped_stream_failure_attribution_reaches_terminal_spend_row() -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + request_start: Final = real_datetime.datetime(2026, 1, 1) + model: Final = "deployment-model" + logging_obj: Final = Logging( + model=model, + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-attribution", + function_id="fn-lit145-attribution", + ) + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-id"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": request_start, + } + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model=model, + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-attribution", + created=1742056047, + model=model, + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="The"))], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + exception: Final = httpx.ReadError("connection dropped") + + with pytest.raises(MidStreamFallbackError): + wrapper._log_stream_failure_and_raise(exception) + await wrapper.wait_for_failure_logging(timeout=10) + + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=exception, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "deployment-id" + assert row["custom_llm_provider"] == "openai" + assert row["model"] == model + + +class _ReplaceSharedFailurePayload(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logging_obj: Logging | None = None + self.replaced: asyncio.Event = asyncio.Event() + self.started: asyncio.Event = asyncio.Event() + self.release: asyncio.Event = asyncio.Event() + self.pause_until_released: bool = False + self.replacement_payload: object | None = None + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: real_datetime.datetime, + end_time: real_datetime.datetime, + ) -> None: + from pydantic import TypeAdapter + + from litellm.types.utils import StandardLoggingPayload + + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + self.started.set() + if self.pause_until_released: + await self.release.wait() + logging_obj: Final = self.logging_obj + if logging_obj is None: + return + replacement: Final = self.replacement_payload or TypeAdapter(StandardLoggingPayload).validate_python( + { + **payload, + "model_id": "newer-deployment-id", + "custom_llm_provider": "newer-provider", + "model": "newer-model", + } + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "standard_logging_object": replacement, + } + self.replaced.set() + + +class _OrderedAttemptFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.attempt_a_started: asyncio.Event = asyncio.Event() + self.attempt_b_started: asyncio.Event = asyncio.Event() + self.release_attempt_a: asyncio.Event = asyncio.Event() + self.release_attempt_b: asyncio.Event = asyncio.Event() + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: real_datetime.datetime, + end_time: real_datetime.datetime, + ) -> None: + from pydantic import TypeAdapter + + from litellm.types.utils import StandardLoggingPayload + + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + if payload["model_id"] == "deployment-a": + self.attempt_a_started.set() + await self.release_attempt_a.wait() + return + if payload["model_id"] == "deployment-b": + self.attempt_b_started.set() + await self.release_attempt_b.wait() + + +class _PausedBaselineFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.started: asyncio.Event = asyncio.Event() + self.release: asyncio.Event = asyncio.Event() + self.standard_logging_object: object | None = None + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: real_datetime.datetime, + end_time: real_datetime.datetime, + ) -> None: + self.standard_logging_object = kwargs.get("standard_logging_object") + self.started.set() + await self.release.wait() + + +@pytest.mark.asyncio +async def test_latest_dropped_attempt_publishes_its_payload_after_earlier_dispatch_finishes() -> None: + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Usage + + request_start: Final = real_datetime.datetime(2026, 1, 1) + marker_a: Final = request_start + marker_b: Final = request_start + real_datetime.timedelta(seconds=1) + recorder: Final = _OrderedAttemptFailureRecorder() + logging_obj: Final = Logging( + model="deployment-model-a", + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-attempt-order", + function_id="fn-lit145-attempt-order", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="deployment-model-a", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-a"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": marker_a, + "combined_usage_object": Usage(prompt_tokens=4, completion_tokens=1, total_tokens=5), + "response_cost": 0.001, + } + + attempt_a: Final = asyncio.create_task( + logging_obj.dispatch_attempt_failure_handlers(httpx.ReadError("attempt A dropped"), "") + ) + await asyncio.wait_for(recorder.attempt_a_started.wait(), timeout=10) + + logging_obj.carry_failed_attempt_usage() + logging_obj.update_environment_variables( + model="deployment-model-b", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-b"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "model": "deployment-model-b", + "custom_llm_provider": "anthropic", + "current_attempt_api_call_start_time": marker_b, + "combined_usage_object": Usage(prompt_tokens=7, completion_tokens=2, total_tokens=9), + "response_cost": 0.002, + } + + attempt_b: Final = asyncio.create_task( + logging_obj.dispatch_attempt_failure_handlers(httpx.ReadError("attempt B dropped"), "") + ) + await asyncio.wait_for(recorder.attempt_b_started.wait(), timeout=10) + recorder.release_attempt_a.set() + await attempt_a + recorder.release_attempt_b.set() + await attempt_b + + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=httpx.ReadError("attempt B dropped"), + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "deployment-b" + assert row["custom_llm_provider"] == "anthropic" + assert row["model"] == "deployment-model-b" + + +@pytest.mark.asyncio +async def test_dropped_attempt_attribution_is_published_while_failure_callback_is_paused() -> None: + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Usage + + request_start: Final = real_datetime.datetime(2026, 1, 1) + marker_a: Final = request_start + failure: Final = httpx.ReadError("attempt A dropped") + recorder: Final = _OrderedAttemptFailureRecorder() + logging_obj: Final = Logging( + model="deployment-model-a", + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-attempt-paused", + function_id="fn-lit145-attempt-paused", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="deployment-model-a", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-a"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": marker_a, + "combined_usage_object": Usage(prompt_tokens=4, completion_tokens=1, total_tokens=5), + "response_cost": 0.001, + } + + attempt: Final = asyncio.create_task(logging_obj.dispatch_attempt_failure_handlers(failure, "")) + await asyncio.wait_for(recorder.attempt_a_started.wait(), timeout=10) + + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=failure, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + try: + assert row["model_id"] == "deployment-a" + assert row["custom_llm_provider"] == "openai" + assert row["model"] == "deployment-model-a" + assert row["prompt_tokens"] == 4 + assert row["completion_tokens"] == 1 + finally: + recorder.release_attempt_a.set() + await attempt + + +@pytest.mark.asyncio +async def test_failed_attempt_publishes_baseline_reason_after_invalidation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pydantic import TypeAdapter + + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.hooks.autorouter_baseline_cache import AutoRouterBaselineCache, CapturedBaselineObservation + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.router import Router + from litellm.types.utils import CallTypes, StandardLoggingPayload, Usage + from tests.unit.proxy.hooks.test_autorouter_baseline_cache import _MODELS, _kwargs + + request_start: Final = real_datetime.datetime(2026, 1, 1) + marker: Final = request_start + model: Final = "deployment-model-a" + router: Final = Router(model_list=_MODELS, disable_cooldowns=True) + baseline_cache: Final = AutoRouterBaselineCache( + None, + router=lambda: router, + clock=lambda: request_start.timestamp() + 10, + ) + recorder: Final = _PausedBaselineFailureRecorder() + logging_obj: Final = Logging( + model=model, + messages=ONE_USER_MESSAGE, + stream=True, + call_type=CallTypes.anthropic_messages.value, + start_time=request_start, + litellm_call_id="call-lit145-baseline-failure", + function_id="fn-lit145-baseline-failure", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-a"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": marker, + "combined_usage_object": Usage(prompt_tokens=4, completion_tokens=1, total_tokens=5), + "response_cost": 0.001, + } + await baseline_cache.async_pre_call_deployment_hook(_kwargs(logging_obj), CallTypes.anthropic_messages) + assert logging_obj.baseline_cache_context is not None + + failure: Final = httpx.ReadError("attempt A dropped") + attempt: Final = asyncio.create_task(logging_obj.dispatch_attempt_failure_handlers(failure, "")) + try: + await asyncio.wait_for(recorder.started.wait(), timeout=10) + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=failure, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "deployment-a" + assert row["custom_llm_provider"] == "openai" + assert row["model"] == model + assert row["prompt_tokens"] == 4 + assert row["completion_tokens"] == 1 + lifted_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(lifted["standard_logging_object"]) + callback_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python( + recorder.standard_logging_object + ) + + def baseline_reason(payload: StandardLoggingPayload) -> str: + encoded: Final = payload["autorouter_baseline_observation"] + assert isinstance(encoded, str) + return CapturedBaselineObservation.model_validate_json(encoded).observation.reason or "" + + assert baseline_reason(lifted_payload) == "failed_request" + assert baseline_reason(callback_payload) == "failed_request" + finally: + recorder.release.set() + await attempt + + +@pytest.mark.asyncio +async def test_queued_stale_attempt_dispatch_does_not_publish_over_current_attempt() -> None: + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Usage + + request_start: Final = real_datetime.datetime(2026, 1, 1) + marker_a: Final = request_start + marker_b: Final = request_start + real_datetime.timedelta(seconds=1) + recorder: Final = _OrderedAttemptFailureRecorder() + logging_obj: Final = Logging( + model="deployment-model-a", + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-stale-attempt", + function_id="fn-lit145-stale-attempt", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="deployment-model-a", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-a"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": marker_a, + "combined_usage_object": Usage(prompt_tokens=4, completion_tokens=1, total_tokens=5), + "response_cost": 0.001, + } + failure_a: Final = httpx.ReadError("attempt A dropped") + coro_a: Final = logging_obj.dispatch_attempt_failure_handlers(failure_a, "") + + logging_obj.carry_failed_attempt_usage() + logging_obj.update_environment_variables( + model="deployment-model-b", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-b"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "model": "deployment-model-b", + "custom_llm_provider": "anthropic", + "current_attempt_api_call_start_time": marker_b, + "combined_usage_object": Usage(prompt_tokens=7, completion_tokens=2, total_tokens=9), + "response_cost": 0.002, + } + failure_b: Final = httpx.ReadError("attempt B dropped") + coro_b: Final = logging_obj.dispatch_attempt_failure_handlers(failure_b, "") + + task_a: Final = asyncio.create_task(coro_a) + task_b: Final = asyncio.create_task(coro_b) + try: + await asyncio.wait_for( + asyncio.gather(recorder.attempt_a_started.wait(), recorder.attempt_b_started.wait()), + timeout=10, + ) + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=failure_b, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "deployment-b" + assert row["custom_llm_provider"] == "anthropic" + assert row["model"] == "deployment-model-b" + finally: + recorder.release_attempt_a.set() + recorder.release_attempt_b.set() + await asyncio.gather(task_a, task_b) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queued_start", [False, True], ids=["running-start", "queued-start"]) +async def test_dropped_stream_failure_does_not_overwrite_newer_terminal_payload(queued_start: bool) -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Delta, ModelResponseStream, StandardLoggingPayload, StreamingChoices, Usage + from pydantic import TypeAdapter + + request_start: Final = real_datetime.datetime(2026, 1, 1) + model: Final = "deployment-model" + recorder: Final = _ReplaceSharedFailurePayload() + logging_obj: Final = Logging( + model=model, + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-identity", + function_id="fn-lit145-identity", + dynamic_async_failure_callbacks=[recorder], + ) + recorder.logging_obj = logging_obj + initial_payload: Final = { + "model_id": "older-deployment-id", + "custom_llm_provider": "older-provider", + "model": "older-model", + } + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-id"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": request_start, + "standard_logging_object": initial_payload, + } + + def assert_newer_payload_attribution() -> None: + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=exception, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "newer-deployment-id" + assert row["custom_llm_provider"] == "newer-provider" + assert row["model"] == "newer-model" + + exception: Final = httpx.ReadError("connection dropped") + if queued_start: + dispatch: Final = logging_obj.dispatch_attempt_failure_handlers(exception, "") + logging_obj._failure_handler_helper_fn( + exception=httpx.ReadError("newer failure"), + traceback_exception="", + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + replacement_base: Final = logging_obj.model_call_details["standard_logging_object"] + replacement_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python( + { + **replacement_base, + "model_id": "newer-deployment-id", + "custom_llm_provider": "newer-provider", + "model": "newer-model", + } + ) + recorder.pause_until_released = True + recorder.replacement_payload = replacement_payload + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "standard_logging_object": replacement_payload, + } + assert logging_obj.model_call_details["current_attempt_api_call_start_time"] == request_start + + task: Final = asyncio.create_task(dispatch) + try: + await asyncio.wait_for(recorder.started.wait(), timeout=10) + assert_newer_payload_attribution() + finally: + recorder.release.set() + await task + + await asyncio.wait_for(recorder.replaced.wait(), timeout=10) + assert_newer_payload_attribution() + return + + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model=model, + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-identity", + created=1742056047, + model=model, + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="The"))], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + + with pytest.raises(MidStreamFallbackError): + wrapper._log_stream_failure_and_raise(exception) + await wrapper.wait_for_failure_logging(timeout=10) + await asyncio.wait_for(recorder.replaced.wait(), timeout=10) + assert_newer_payload_attribution() + + def test_a_timed_out_attempt_from_the_group_a_fallback_left_is_still_billed(): """After a fallback hop the router restarts ``attempted_retries`` at 0 for the new model group and re-logs the attempt that exhausted the previous group, so the final attempt (dep-b, retry 1) must be @@ -2655,6 +3372,124 @@ def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rat assert response_cost == pytest.approx(expected_input_cost) + +COMPACTION_INSTRUCTIONS = {"role": "system", "content": "answer briefly"} +COMPACTED_HISTORY = [ + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +] +COMPACTION_ITEM = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +TURN_AFTER_COMPACTION = {"role": "user", "content": "what is the cat's name?"} + + +def _failed_call_prompt_tokens(call_type, dispatched_call_type, request_input): + from litellm.proxy.utils import _failure_usage_to_lift + + lifted = _failure_usage_to_lift( + model_call_details={"call_type": call_type, "model": FAILURE_USAGE_MODEL, "messages": request_input}, + request_body={}, + dispatched=True, + dispatched_call_type=dispatched_call_type, + ) + assert lifted is not None + return lifted[0].prompt_tokens + + +@pytest.mark.parametrize( + "call_type, dispatched_call_type", + [ + ("acompletion", "acompletion"), + ("aanthropic_messages", "anthropic_messages"), + ("responses", "acompletion"), + ], +) +def test_a_failed_call_off_the_responses_route_counts_the_history_before_a_compaction_typed_message( + call_type, dispatched_call_type +): + def prompt_tokens(request_input): + return _failed_call_prompt_tokens(call_type, dispatched_call_type, request_input) + + sent = prompt_tokens([COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + assert sent > prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + +def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest_compaction_item(): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.utils import _failure_fields_to_lift + + def earlier_attempt_prompt_tokens(request_input): + logging_obj = _dispatched_failure_logging_obj( + FAILURE_USAGE_MODEL, + get_litellm_params( + metadata={ + "attempted_retries": 1, + "model_info": {"id": "dep-compaction-final"}, + "previous_models": ( + {"deployment_id": "dep-compaction-final", "exception_type": "Timeout", "attempted_retries": 0}, + ), + } + ), + ) + logging_obj.model_call_details["call_type"] = "aresponses" + logging_obj.model_call_details["messages"] = request_input + lifted = _failure_fields_to_lift( + {"litellm_logging_obj": logging_obj}, provider_billed=True, dispatched_call_type="aresponses" + ) + (earlier_attempt,) = lifted["billed_failed_attempts"] + return earlier_attempt.usage.prompt_tokens + + sent = earlier_attempt_prompt_tokens( + [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + + assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + +@pytest.mark.parametrize( + "current_attempt_api_call_start_time", + [None, 1700000001.0], + ids=["retry-not-dispatched", "retry-dispatched"], +) +def test_a_retry_failure_bills_input_only_for_dispatched_attempts( + current_attempt_api_call_start_time: float | None, +) -> None: + from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_fields_to_lift + from litellm.types.utils import Usage + + prior_usage: Final = Usage(prompt_tokens=4, completion_tokens=2, total_tokens=6) + estimated_retry_usage: Final = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) + assert estimated_retry_usage is not None + + lifted: Final = _failure_fields_to_lift( + { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "current_attempt_api_call_start_time": current_attempt_api_call_start_time, + "failed_attempts_usage": ((prior_usage, 0.25),), + "call_type": "acompletion", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + } + ) + } + ) + + usage: Final = lifted["combined_usage_object"] + assert isinstance(usage, Usage) + expected_retry_tokens: Final = ( + estimated_retry_usage.prompt_tokens if current_attempt_api_call_start_time is not None else 0 + ) + assert usage.prompt_tokens == prior_usage.prompt_tokens + expected_retry_tokens + assert usage.completion_tokens == prior_usage.completion_tokens + assert usage.total_tokens == prior_usage.total_tokens + expected_retry_tokens + assert lifted["response_cost"] == 0.25 + + @pytest.mark.parametrize( "model_call_details, dispatched", [ diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 8928b79bcd4..6044475652c 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,19 +1,42 @@ +import asyncio import contextlib +import gc +import threading +import weakref import httpx import json import litellm import pytest +import redis import respx +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping +from datetime import datetime from types import SimpleNamespace +from uuid import uuid4 from typing import Final from unittest.mock import MagicMock +from fakeredis import TcpFakeServer from fastapi import HTTPException -from litellm.proxy.route_llm_request import JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, route_request -from litellm.types.router import RouterGeneralSettings +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.route_llm_request import ( + JSON_OBJECT_ADAPTER, + ProxyModelNotFoundError, + UserConfigRouters, + route_request, + user_config_routers, +) +from litellm.proxy.utils import ProxyLogging +from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ADAPTIVE_ROUTER_RESPONSE_HEADER, +) +from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError @pytest.mark.parametrize( @@ -2434,3 +2457,571 @@ async def test_team_key_search_of_registered_bedrock_store_reaches_bedrock_with_ posted_url ) assert "Credential=AKIAREGISTRYKEY00000/" in signed_auth, signed_auth + + +@pytest.fixture(autouse=True) +def discard_cached_user_config_routers() -> Iterator[None]: + yield + user_config_routers.clear() + + +@pytest.fixture +def fake_redis_port() -> Iterator[int]: + server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis") + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield server.server_address[1] + server.shutdown() + server.server_close() + + +async def _wait_out_a_rate_limit_minute_rollover() -> None: + seconds_left: Final = 60 - datetime.now().second + if seconds_left < 10: + await asyncio.sleep(seconds_left) + + +def _proxy_logging_obj(model: str = "rate-limited") -> LiteLLMLoggingObj: + logging_obj, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model=model, + messages=[{"role": "user", "content": "hi"}], + litellm_call_id=str(uuid4()), + ) + assert isinstance(logging_obj, LiteLLMLoggingObj) + return logging_obj + + +def _mock_deployment(model_name: str, **deployment_limits: int) -> dict[str, object]: + return { + "model_name": model_name, + **deployment_limits, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "mock_response": "ok"}, + } + + +async def _user_config_request_outcome( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + logging_obj: LiteLLMLoggingObj | None, +) -> str: + data: Final = { + "model": "rate-limited", + "messages": [{"role": "user", "content": "hi"}], + "num_retries": 0, + **({"litellm_logging_obj": logging_obj} if logging_obj is not None else {}), + "user_config": { + **routing_settings, + **({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}), + "model_list": [_mock_deployment("rate-limited", **deployment_limits)], + }, + } + try: + await (await route_request(data, None, None, "acompletion")) + except (litellm.RateLimitError, RouterRateLimitError): + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + +async def _tpm_for_three_requests() -> int: + messages: Final = [{"role": "user", "content": "hi"}] + probe: Final = await litellm.acompletion(model="openai/gpt-5.5", messages=messages, mock_response="ok") + assert isinstance(probe, litellm.ModelResponse) + probe_usage: Final = getattr(probe, "usage", None) + assert isinstance(probe_usage, litellm.Usage) + return 2 * probe_usage.total_tokens + litellm.token_counter(messages=messages) + + +async def _request_cost() -> float: + probe: Final = await litellm.acompletion( + model="openai/gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok" + ) + return litellm.completion_cost(completion_response=probe) + + +async def _proxy_router_outcomes( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + requests: int, +) -> list[str]: + router: Final = litellm.Router( + **routing_settings, + num_retries=0, + model_list=[_mock_deployment("proxy-rate-limited", **deployment_limits)], + **({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}), + ) + + async def outcome() -> str: + try: + await router.acompletion(model="proxy-rate-limited", messages=[{"role": "user", "content": "hi"}]) + except (litellm.RateLimitError, RouterRateLimitError): + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + try: + return [await outcome() for _ in range(requests)] + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() -> None: + data: Final = { + "model": "smart-cheap-router", + "messages": [{"role": "user", "content": "Write a one line python function that adds two numbers"}], + "metadata": {}, + "user_config": { + "model_list": [ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast", "smart"]}, + }, + }, + { + **_mock_deployment("fast"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 2, "strengths": []}}, + }, + { + **_mock_deployment("smart"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}}, + }, + ] + }, + } + + await (await route_request(data, None, None, "acompletion")) + headers: Final = await ProxyLogging(user_api_key_cache=DualCache()).post_call_response_headers_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), response=None + ) + + assert headers == {ADAPTIVE_ROUTER_RESPONSE_HEADER: data["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY]} + + +def _global_router_callbacks() -> tuple[list[object], list[object]]: + return list(litellm.callbacks), list(litellm.input_callback) + + +@pytest.mark.parametrize( + "routing_settings, deployment_limits", + [ + ({"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}), + ({"routing_strategy": "usage-based-routing-v2"}, {"tpm": 30}), + ({"routing_strategy": "usage-based-routing"}, {"rpm": 2}), + ({"optional_pre_call_checks": ["enforce_model_rate_limits"]}, {"rpm": 2}), + ], +) +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.parametrize("with_redis", [True, False]) +@pytest.mark.asyncio +async def test_user_config_router_enforces_its_limit_across_proxy_requests( + fake_redis_port: int, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + arrives_with_logging_obj: bool, + with_redis: bool, +) -> None: + await _wait_out_a_rate_limit_minute_rollover() + callbacks_before: Final = _global_router_callbacks() + outcomes: Final = [ + await _user_config_request_outcome( + fake_redis_port if with_redis else None, + routing_settings, + deployment_limits, + _proxy_logging_obj() if arrives_with_logging_obj else None, + ) + for _ in range(6) + ] + + assert "429" in outcomes, outcomes + assert _global_router_callbacks() == callbacks_before + assert outcomes == await _proxy_router_outcomes( + fake_redis_port if with_redis else None, routing_settings, deployment_limits, len(outcomes) + ) + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_limit( + fake_redis_port: int, arrives_with_logging_obj: bool +) -> None: + tpm_for_three_requests: Final = await _tpm_for_three_requests() + + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, + {"routing_strategy": "usage-based-routing-v2"}, + {"tpm": tpm_for_three_requests}, + _proxy_logging_obj() if arrives_with_logging_obj else None, + ) + + await _wait_out_a_rate_limit_minute_rollover() + concurrent_outcomes: Final = await asyncio.gather(request(), request()) + later_outcomes: Final = [await request() for _ in range(2)] + + assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"] + + +@pytest.mark.asyncio +async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None: + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}, _proxy_logging_obj() + ) + + await _wait_out_a_rate_limit_minute_rollover() + outcomes: Final = await asyncio.gather(*(request() for _ in range(4))) + + assert outcomes.count("200") <= 2, outcomes + assert "429" in outcomes, outcomes + + +@pytest.mark.parametrize( + "routing_strategy, limited_by", + [("usage-based-routing-v2", "rpm"), ("usage-based-routing-v2", "tpm"), ("usage-based-routing", "rpm")], +) +@pytest.mark.parametrize("proxy_router_shares_redis", [True, False]) +@pytest.mark.asyncio +async def test_a_proxy_router_on_the_same_strategy_leaves_the_user_config_limit_unchanged( + fake_redis_port: int, routing_strategy: str, limited_by: str, proxy_router_shares_redis: bool +) -> None: + deployment_limits: Final = {"rpm": 3} if limited_by == "rpm" else {"tpm": await _tpm_for_three_requests()} + + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, {"routing_strategy": routing_strategy}, deployment_limits, _proxy_logging_obj() + ) + + async def outcomes() -> list[str]: + first_instance: Final = [await request() for _ in range(2)] + user_config_routers.clear() + second_instance: Final = [await request() for _ in range(3)] + return [*first_instance, *second_instance] + + await _wait_out_a_rate_limit_minute_rollover() + alone: Final = await outcomes() + user_config_routers.clear() + await asyncio.sleep(0.3) + redis.Redis(host="127.0.0.1", port=fake_redis_port).flushall() + proxy_router: Final = litellm.Router( + routing_strategy=routing_strategy, + model_list=[_mock_deployment("proxy-model")], + redis_host="127.0.0.1" if proxy_router_shares_redis else None, + redis_port=fake_redis_port if proxy_router_shares_redis else None, + ) + try: + beside_proxy_router: Final = await outcomes() + finally: + proxy_router.discard() + + assert "429" in alone, alone + assert beside_proxy_router == alone + + +@pytest.mark.parametrize("user_config_settings", [{}, {"optional_pre_call_checks": ["router_budget_limiting"]}]) +@pytest.mark.asyncio +async def test_user_config_requests_stay_under_the_proxy_provider_budget( + user_config_settings: Mapping[str, object], +) -> None: + proxy_router: Final = litellm.Router( + model_list=[_mock_deployment("proxy-model")], + provider_budget_config={"openai": {"budget_limit": await _request_cost() / 2, "time_period": "1d"}}, + ) + try: + await proxy_router.acompletion(model="proxy-model", messages=[{"role": "user", "content": "hi"}]) + await asyncio.sleep(0.3) + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await _user_config_request_outcome(None, user_config_settings, {}, _proxy_logging_obj()) + finally: + proxy_router.discard() + + +@pytest.mark.asyncio +async def test_a_user_config_provider_budget_limits_its_own_requests() -> None: + budget_limit: Final = await _request_cost() / 2 + + def request() -> Coroutine[None, None, str]: + budget: Final = {"openai": {"budget_limit": budget_limit, "time_period": "1d"}} + return _user_config_request_outcome(None, {"provider_budget_config": budget}, {}, _proxy_logging_obj()) + + first: Final = await request() + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await request() + assert first == "200" + + +@pytest.mark.asyncio +async def test_a_request_in_flight_when_its_user_config_router_is_dropped_still_counts_against_its_budget( + fake_redis_port: int, +) -> None: + budget_limit: Final = 1.5 * await _request_cost() + + async def request() -> None: + await ( + await route_request( + { + "model": "budgeted", + "messages": [{"role": "user", "content": "hi"}], + "litellm_logging_obj": _proxy_logging_obj("budgeted"), + "user_config": { + "redis_host": "127.0.0.1", + "redis_port": fake_redis_port, + "provider_budget_config": {"openai": {"budget_limit": budget_limit, "time_period": "1d"}}, + "model_list": [ + { + "model_name": "budgeted", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "mock_response": "ok", + "mock_delay": 0.5, + }, + } + ], + }, + }, + None, + None, + "acompletion", + ) + ) + + await request() + in_flight: Final = asyncio.create_task(request()) + await asyncio.sleep(0.2) + user_config_routers.clear() + await in_flight + await asyncio.sleep(0.3) + + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await request() + + +async def _user_config_request(model: str, user_config: Mapping[str, object], arrives_with_logging_obj: bool) -> object: + return await ( + await route_request( + { + "model": model, + "messages": list(_HELLO_MESSAGES), + "num_retries": 0, + **({"litellm_logging_obj": _proxy_logging_obj(model)} if arrives_with_logging_obj else {}), + "user_config": user_config, + }, + None, + None, + "acompletion", + ) + ) + + +def _ordered_pair(model_name: str, first_base: str, second_base: str) -> list[dict[str, object]]: + return [ + { + "model_name": model_name, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "api_base": api_base, "order": order}, + } + for order, api_base in ((1, first_base), (2, second_base)) + ] + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_a_rate_limited_user_config_deployment_cools_down_for_the_next_requests( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + rate_limited_base: Final = "https://rate-limited.example/v1" + healthy_base: Final = "https://healthy.example/v1" + user_config: Final = {"model_list": _ordered_pair("flaky", rate_limited_base, healthy_base)} + + with respx.mock(assert_all_called=False) as respx_mock: + rate_limited: Final = respx_mock.post(f"{rate_limited_base}/chat/completions").mock( + return_value=httpx.Response(429, json={"error": {"message": "slow down"}}) + ) + respx_mock.post(f"{healthy_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5")) + for _ in range(3): + await _user_config_request("flaky", user_config, arrives_with_logging_obj) + await asyncio.sleep(0.3) + + assert rate_limited.call_count == 1 + + +async def _remaining_usage_headers_after_429_then_200s( + send: Callable[[], Awaitable[object]], api_base: str +) -> list[tuple[object, object] | str]: + async def remaining_usage() -> tuple[object, object] | str: + try: + response: Final = await send() + except litellm.RateLimitError: + return "429" + finally: + await asyncio.sleep(0.3) + assert isinstance(response, litellm.ModelResponse) + headers: Final = response._hidden_params["additional_headers"] + return headers.get("x-ratelimit-remaining-requests"), headers.get("x-ratelimit-remaining-tokens") + + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.post(f"{api_base}/chat/completions").mock( + side_effect=[ + httpx.Response(429, json={"error": {"message": "slow down"}}), + _chat_completion_response("gpt-5.5"), + _chat_completion_response("gpt-5.5"), + ] + ) + return [await remaining_usage() for _ in range(3)] + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_user_config_remaining_usage_headers_match_a_proxy_router_with_the_same_deployment( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await _wait_out_a_rate_limit_minute_rollover() + proxy_base: Final = "https://proxy-limited.example/v1" + user_config_base: Final = "https://user-config-limited.example/v1" + + def deployment(model_name: str, api_base: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "api_base": api_base, + "rpm": 10, + "tpm": 1000, + }, + } + + proxy_router: Final = litellm.Router(model_list=[deployment("proxy-limited", proxy_base)]) + try: + proxy_router_headers: Final = await _remaining_usage_headers_after_429_then_200s( + lambda: proxy_router.acompletion(model="proxy-limited", messages=list(_HELLO_MESSAGES), num_retries=0), + proxy_base, + ) + finally: + proxy_router.discard() + + user_config: Final = {"model_list": [deployment("user-config-limited", user_config_base)]} + user_config_headers: Final = await _remaining_usage_headers_after_429_then_200s( + lambda: _user_config_request("user-config-limited", user_config, arrives_with_logging_obj), user_config_base + ) + + assert user_config_headers == proxy_router_headers + + +async def _first_deployment_calls_over_successes_then_failures( + send: Callable[[], Awaitable[object]], first_base: str, second_base: str +) -> int: + successes: Final = [_chat_completion_response("gpt-5.5") for _ in range(6)] + failures: Final = [httpx.Response(500, json={"error": {"message": "boom"}}) for _ in range(6)] + with respx.mock(assert_all_called=False) as respx_mock: + first: Final = respx_mock.post(f"{first_base}/chat/completions").mock(side_effect=[*successes, *failures]) + respx_mock.post(f"{second_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5")) + for _ in range(len(successes) + len(failures)): + await send() + await asyncio.sleep(0.3) + return first.call_count + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_user_config_successes_keep_a_deployment_out_of_cooldown_like_a_proxy_router( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await _wait_out_a_rate_limit_minute_rollover() + proxy_router: Final = litellm.Router( + model_list=_ordered_pair("proxy-pair", "https://proxy-first.example/v1", "https://proxy-second.example/v1") + ) + try: + proxy_router_calls: Final = await _first_deployment_calls_over_successes_then_failures( + lambda: proxy_router.acompletion(model="proxy-pair", messages=list(_HELLO_MESSAGES), num_retries=0), + "https://proxy-first.example/v1", + "https://proxy-second.example/v1", + ) + finally: + proxy_router.discard() + + user_config: Final = { + "model_list": _ordered_pair("user-pair", "https://user-first.example/v1", "https://user-second.example/v1") + } + + user_config_calls: Final = await _first_deployment_calls_over_successes_then_failures( + lambda: _user_config_request("user-pair", user_config, arrives_with_logging_obj), + "https://user-first.example/v1", + "https://user-second.example/v1", + ) + + assert user_config_calls == proxy_router_calls + + +@pytest.mark.asyncio +async def test_user_config_adaptive_router_learns_from_proxy_requests() -> None: + user_config: Final = { + "model_list": [ + _mock_deployment("smart"), + _mock_deployment("fast"), + { + "model_name": "adaptive", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["smart", "fast"]}, + }, + }, + ] + } + await ( + await route_request( + { + "model": "adaptive", + "messages": [{"role": "user", "content": "write a sorting function"}], + "metadata": {"session_id": "session-1"}, + "litellm_logging_obj": _proxy_logging_obj("adaptive"), + "user_config": user_config, + }, + None, + None, + "acompletion", + ) + ) + await asyncio.sleep(0.3) + + router: Final = user_config_routers.get(user_config, lambda: pytest.fail("the user_config router was not reused")) + snapshot: Final = await router.adaptive_routers["adaptive"][0].strategy.get_state_snapshot() + assert snapshot["feedback_contexts_live"] == 1, snapshot + + +@pytest.mark.asyncio +async def test_user_config_routers_discard_the_least_recently_used_router_past_capacity() -> None: + routers: Final = UserConfigRouters(max_size=2) + + def config(model_name: str) -> dict[str, object]: + return {"routing_strategy": "usage-based-routing-v2", "model_list": [_mock_deployment(model_name)]} + + def build(model_name: str) -> litellm.Router: + return litellm.Router(**config(model_name)) + + def must_be_cached() -> litellm.Router: + pytest.fail("a cached router was rebuilt") + + first_handler: Final = weakref.ref(routers.get(config("first"), lambda: build("first")).lowesttpm_logger_v2) + second_handler: Final = weakref.ref(routers.get(config("second"), lambda: build("second")).lowesttpm_logger_v2) + routers.get(config("first"), must_be_cached) + routers.get(config("third"), lambda: build("third")) + await asyncio.sleep(0.2) + gc.collect() + + assert routers.get(config("first"), must_be_cached) is not None + assert first_handler() is not None + assert second_handler() is None + routers.clear() diff --git a/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py index 25f1301dd9d..b5a41c31810 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py @@ -12,10 +12,11 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app from litellm.proxy.ui_crud_endpoints.latest_release_endpoints import ( + CHANGELOG_FEED_URL, + CHANGELOG_URL, LATEST_RELEASE_CACHE_KEY, LATEST_RELEASE_CACHE_TTL_SECONDS, LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS, - LATEST_RELEASE_URL, RELEASE_BY_TAG_URL, LatestReleaseInfo, LatestReleaseUnavailable, @@ -49,37 +50,68 @@ SAMPLE_RELEASE: Final = { "html_url": "https://github.com/BerriAI/litellm/releases/tag/v1.102.0", "body": SAMPLE_BODY, } + + +def _changelog_page(folder: str) -> str: + return f"{CHANGELOG_URL}{folder}/{folder.removeprefix('v').replace('.', '-')}" + + +def _feed(*folders: str) -> httpx.Response: + items: Final = "".join( + f"{folder}{_changelog_page(folder)}" for folder in folders + ) + return httpx.Response(200, content=f'{items}'.encode()) + + EXPECTED_INFO: Final = { "version": "1.102.0", "new_features": 2, "bug_fixes": 2, "other_updates": 4, - "release_url": SAMPLE_RELEASE["html_url"], + "release_url": _changelog_page("v1.102.0"), } -class _RecordingClient: - def __init__(self, outcomes: list[httpx.Response | Exception]) -> None: - self._outcomes = outcomes - self.calls: list[tuple[str, float | None]] = [] - - async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: - self.calls.append((url, timeout)) - outcome = self._outcomes[min(len(self.calls) - 1, len(self._outcomes) - 1)] - if isinstance(outcome, Exception): - raise outcome - return outcome - - def _github_response(status: int = 200, payload: object = SAMPLE_RELEASE) -> httpx.Response: return httpx.Response(status, content=json.dumps(payload).encode()) +def _tag_url(tag: str) -> str: + return RELEASE_BY_TAG_URL.format(tag=tag) + + +class _ByUrl: + """Answers each URL from a fixed map; unknown URLs are 404, like a tag with no release.""" + + def __init__(self, routes: dict[str, httpx.Response | Exception]) -> None: + self._routes = routes + self.calls: list[tuple[str, float | None]] = [] + + async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + self.calls.append((url, timeout)) + outcome = self._routes.get(url, httpx.Response(404, content=b"{}")) + if isinstance(outcome, Exception): + raise outcome + return outcome + + @property + def urls(self) -> list[str]: + return [url for url, _ in self.calls] + + +def _minor_release_routes(**overrides: httpx.Response | Exception) -> dict[str, httpx.Response | Exception]: + return {CHANGELOG_FEED_URL: _feed("v1.102.0"), _tag_url("v1.102.0"): _github_response(), **overrides} + + +def _minor_release(**overrides: httpx.Response | Exception) -> _ByUrl: + return _ByUrl(_minor_release_routes(**overrides)) + + def _fresh_cache() -> InMemoryCache: return InMemoryCache(max_size_in_memory=1, default_ttl=LATEST_RELEASE_CACHE_TTL_SECONDS) -def _override_dependencies(client: _RecordingClient, cache: InMemoryCache, role: LitellmUserRoles) -> None: +def _override_dependencies(client: _ByUrl, cache: InMemoryCache, role: LitellmUserRoles) -> None: async def auth() -> UserAPIKeyAuth: return UserAPIKeyAuth(api_key="sk-test", user_id="test-user", user_role=role) @@ -127,35 +159,68 @@ class TestCountReleaseBullets: class TestGetLatestReleaseInfo: @pytest.mark.asyncio - async def test_fetches_and_parses_github_release(self): - client = _RecordingClient([_github_response()]) + async def test_announces_the_changelog_release_and_links_its_changelog_page(self): + client = _minor_release() result = await get_latest_release_info(client=client, cache=_fresh_cache(), fetch_lock=asyncio.Lock()) assert isinstance(result, LatestReleaseInfo) assert result.model_dump() == EXPECTED_INFO - assert client.calls == [(LATEST_RELEASE_URL, 5)] + assert client.calls == [(CHANGELOG_FEED_URL, 5), (_tag_url("v1.102.0"), 5)] + + @pytest.mark.asyncio + async def test_github_release_missing_from_the_changelog_is_not_announced(self): + automated: Final = _github_response(payload=_release("v1.103.0", "* feat(proxy): automated release\n")) + result = await _info(_minor_release(**{_tag_url("v1.103.0"): automated})) + assert isinstance(result, LatestReleaseInfo) + assert result.model_dump() == EXPECTED_INFO + + @pytest.mark.asyncio + async def test_highest_stable_version_wins_regardless_of_feed_order(self): + client = _minor_release(**{CHANGELOG_FEED_URL: _feed("v1.101.6", "v1.103.0rc1", "v1.102.0", "v1.100.9")}) + result = await _info(client) + assert isinstance(result, LatestReleaseInfo) + assert result.model_dump() == EXPECTED_INFO + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "feed", + [ + httpx.Response(503, content=_feed("v1.102.0").content), + httpx.Response(200, content=b"not a feed"), + _feed("v1.103.0rc1", "v1.103.0-dev.1"), + httpx.Response( + 200, + content=f"https://example.com/?{_changelog_page('v9.0.0')}" + "https://example.com/release_notes/v9.0.0/x".encode(), + ), + ], + ids=["server_error", "not_xml", "only_prereleases", "foreign_links"], + ) + async def test_bad_changelog_feeds_are_unavailable(self, feed: httpx.Response): + client = _minor_release(**{CHANGELOG_FEED_URL: feed}) + result = await _info(client) + assert isinstance(result, LatestReleaseUnavailable) + assert client.urls == [CHANGELOG_FEED_URL] @pytest.mark.asyncio async def test_second_call_within_ttl_does_not_refetch(self): - client = _RecordingClient([_github_response()]) + client = _minor_release() cache = _fresh_cache() fetch_lock = asyncio.Lock() first = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) second = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) assert first == second - assert len(client.calls) == 1 + assert len(client.calls) == 2 @pytest.mark.asyncio async def test_success_is_cached_for_the_full_ttl(self): cache = _fresh_cache() - await get_latest_release_info( - client=_RecordingClient([_github_response()]), cache=cache, fetch_lock=asyncio.Lock() - ) + await get_latest_release_info(client=_minor_release(), cache=cache, fetch_lock=asyncio.Lock()) remaining = await cache.async_get_ttl(LATEST_RELEASE_CACHE_KEY) - time.time() assert LATEST_RELEASE_CACHE_TTL_SECONDS - 5 < remaining <= LATEST_RELEASE_CACHE_TTL_SECONDS @pytest.mark.asyncio async def test_failure_is_cached_briefly_so_github_is_not_hammered(self): - client = _RecordingClient([httpx.ConnectError("boom")]) + client = _ByUrl({CHANGELOG_FEED_URL: httpx.ConnectError("boom")}) cache = _fresh_cache() fetch_lock = asyncio.Lock() first = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) @@ -172,30 +237,29 @@ class TestGetLatestReleaseInfo: @pytest.mark.parametrize( "response", [ + _github_response(status=404, payload={"message": "Not Found"}), _github_response(status=403, payload={"message": "rate limited"}), _github_response(status=500, payload={}), _github_response(payload={"tag_name": "v1.0.0"}), httpx.Response(200, content=b"not json"), + httpx.ConnectError("boom"), ], - ids=["rate_limited", "server_error", "missing_fields", "not_json"], + ids=["no_release_for_tag", "rate_limited", "server_error", "missing_fields", "not_json", "unreachable"], ) - async def test_bad_github_responses_are_unavailable(self, response: httpx.Response): - result = await get_latest_release_info( - client=_RecordingClient([response]), cache=_fresh_cache(), fetch_lock=asyncio.Lock() - ) + async def test_bad_github_responses_are_unavailable(self, response: httpx.Response | Exception): + result = await _info(_minor_release(**{_tag_url("v1.102.0"): response})) assert isinstance(result, LatestReleaseUnavailable) @pytest.mark.asyncio async def test_concurrent_misses_share_one_fetch(self): event = asyncio.Event() - class _BlockingClient(_RecordingClient): + class _BlockingClient(_ByUrl): async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: - self.calls.append((url, timeout)) await event.wait() - return _github_response() + return await super().get(url, timeout=timeout) - client = _BlockingClient([]) + client = _BlockingClient(_minor_release_routes()) cache = _fresh_cache() fetch_lock = asyncio.Lock() tasks = [ @@ -208,19 +272,19 @@ class TestGetLatestReleaseInfo: results = await asyncio.gather(*tasks) expected: Final = LatestReleaseInfo.model_validate(EXPECTED_INFO) assert results == [expected] * 5 - assert len(client.calls) == 1 + assert len(client.calls) == 2 @pytest.mark.asyncio async def test_failure_under_lock_is_also_coalesced(self): event = asyncio.Event() - class _FailingBlockingClient(_RecordingClient): + class _FailingBlockingClient(_ByUrl): async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: self.calls.append((url, timeout)) await event.wait() raise httpx.ConnectError("boom") - client = _FailingBlockingClient([]) + client = _FailingBlockingClient({}) cache = _fresh_cache() fetch_lock = asyncio.Lock() tasks = [ @@ -245,40 +309,23 @@ def _release(tag: str, body: str | None, *, prerelease: bool = False) -> dict[st } -def _tag_url(tag: str) -> str: - return RELEASE_BY_TAG_URL.format(tag=tag) - - -class _GitHubByUrl: - """Answers each GitHub URL from a fixed map; unknown URLs are 404, like a tag with no release.""" - - def __init__(self, routes: dict[str, httpx.Response | Exception]) -> None: - self._routes = routes - self.urls: list[str] = [] - - async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: - self.urls.append(url) - outcome = self._routes.get(url, httpx.Response(404, content=b"{}")) - if isinstance(outcome, Exception): - raise outcome - return outcome - - PATCH_RELEASE: Final = _release("v1.102.2", "* fix(proxy): second backport by @k in https://x/pull/10\n") FIRST_PATCH: Final = _release("v1.102.1", "* fix(proxy): first backport by @k in https://x/pull/9\n") -def _minor_line(**overrides: httpx.Response | Exception) -> dict[str, httpx.Response | Exception]: - routes: dict[str, httpx.Response | Exception] = { - LATEST_RELEASE_URL: _github_response(payload=PATCH_RELEASE), - _tag_url("v1.102.0"): _github_response(payload=SAMPLE_RELEASE), - _tag_url("v1.102.1"): _github_response(payload=FIRST_PATCH), - } - routes.update({_tag_url(tag): outcome for tag, outcome in overrides.items()}) - return routes +def _minor_line(**overrides: httpx.Response | Exception) -> _ByUrl: + return _ByUrl( + { + CHANGELOG_FEED_URL: _feed("v1.102.2", "v1.101.6"), + _tag_url("v1.102.2"): _github_response(payload=PATCH_RELEASE), + _tag_url("v1.102.0"): _github_response(payload=SAMPLE_RELEASE), + _tag_url("v1.102.1"): _github_response(payload=FIRST_PATCH), + **{_tag_url(tag): outcome for tag, outcome in overrides.items()}, + } + ) -async def _info(client: _GitHubByUrl) -> LatestReleaseInfo | LatestReleaseUnavailable: +async def _info(client: _ByUrl) -> LatestReleaseInfo | LatestReleaseUnavailable: return await get_latest_release_info(client=client, cache=_fresh_cache(), fetch_lock=asyncio.Lock()) @@ -289,7 +336,7 @@ class TestPatchReleaseAggregatesMinorLine: @pytest.mark.asyncio async def test_patch_release_counts_include_every_earlier_patch_on_its_line(self): - client = _GitHubByUrl(_minor_line()) + client = _minor_line() result = await _info(client) assert isinstance(result, LatestReleaseInfo) assert result.model_dump() == { @@ -297,23 +344,11 @@ class TestPatchReleaseAggregatesMinorLine: "new_features": 2, "bug_fixes": 4, "other_updates": 4, - "release_url": PATCH_RELEASE["html_url"], + "release_url": _changelog_page("v1.102.2"), } - assert sorted(client.urls) == sorted([LATEST_RELEASE_URL, _tag_url("v1.102.0"), _tag_url("v1.102.1")]) - - @pytest.mark.asyncio - async def test_minor_release_makes_a_single_request(self): - client = _RecordingClient([_github_response()]) - await get_latest_release_info(client=client, cache=_fresh_cache(), fetch_lock=asyncio.Lock()) - assert client.calls == [(LATEST_RELEASE_URL, 5)] - - @pytest.mark.asyncio - async def test_non_semver_latest_tag_counts_only_its_own_notes(self): - client = _GitHubByUrl({LATEST_RELEASE_URL: _github_response(payload=_release("v1.102.1-hotfix", "* fix: x\n"))}) - result = await _info(client) - assert isinstance(result, LatestReleaseInfo) - assert (result.new_features, result.bug_fixes, result.other_updates) == (0, 1, 0) - assert client.urls == [LATEST_RELEASE_URL] + assert sorted(client.urls) == sorted( + [CHANGELOG_FEED_URL, _tag_url("v1.102.2"), _tag_url("v1.102.0"), _tag_url("v1.102.1")] + ) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -326,15 +361,13 @@ class TestPatchReleaseAggregatesMinorLine: ids=["never_published", "prerelease", "null_body"], ) async def test_missing_prerelease_or_empty_patch_adds_nothing(self, skipped: httpx.Response): - result = await _info(_GitHubByUrl(_minor_line(**{"v1.102.1": skipped}))) + result = await _info(_minor_line(**{"v1.102.1": skipped})) assert isinstance(result, LatestReleaseInfo) assert (result.new_features, result.bug_fixes, result.other_updates) == (2, 3, 4) @pytest.mark.asyncio async def test_latest_release_with_null_body_still_counts_its_line(self): - routes = _minor_line() - routes[LATEST_RELEASE_URL] = _github_response(payload=_release("v1.102.2", None)) - result = await _info(_GitHubByUrl(routes)) + result = await _info(_minor_line(**{"v1.102.2": _github_response(payload=_release("v1.102.2", None))})) assert isinstance(result, LatestReleaseInfo) assert (result.new_features, result.bug_fixes, result.other_updates) == (2, 3, 4) @@ -349,31 +382,51 @@ class TestPatchReleaseAggregatesMinorLine: ids=["unreachable", "rate_limited", "not_json"], ) async def test_earlier_patch_failure_is_unavailable(self, outcome: httpx.Response | Exception): - result = await _info(_GitHubByUrl(_minor_line(**{"v1.102.0": outcome}))) + result = await _info(_minor_line(**{"v1.102.0": outcome})) assert isinstance(result, LatestReleaseUnavailable) class TestLatestReleaseInfoEndpoint: def test_returns_release_stats_for_authenticated_user(self, http_client): - _override_dependencies(_RecordingClient([_github_response()]), _fresh_cache(), LitellmUserRoles.INTERNAL_USER) + _override_dependencies(_minor_release(), _fresh_cache(), LitellmUserRoles.INTERNAL_USER) response = http_client.get("/get/latest_release_info") assert response.status_code == 200 assert response.json() == EXPECTED_INFO - def test_returns_null_when_github_is_unreachable(self, http_client): + def test_returns_null_when_the_changelog_is_unreachable(self, http_client): _override_dependencies( - _RecordingClient([httpx.ConnectError("boom")]), _fresh_cache(), LitellmUserRoles.PROXY_ADMIN + _ByUrl({CHANGELOG_FEED_URL: httpx.ConnectError("boom")}), _fresh_cache(), LitellmUserRoles.PROXY_ADMIN ) response = http_client.get("/get/latest_release_info") assert response.status_code == 200 assert response.json() is None + @pytest.mark.parametrize("name", ["LITELLM_DISABLE_LATEST_RELEASE_CHECK", "LITELLM_OFFLINE_MODE"]) + @pytest.mark.parametrize( + ("flag", "expected_body", "expected_urls"), + [ + ("true", None, []), + (" TRUE ", None, []), + ("false", EXPECTED_INFO, [CHANGELOG_FEED_URL, _tag_url("v1.102.0")]), + ], + ) + def test_disable_flag_skips_every_outbound_request( + self, http_client, monkeypatch, name, flag, expected_body, expected_urls + ): + monkeypatch.setenv(name, flag) + client = _minor_release() + _override_dependencies(client, _fresh_cache(), LitellmUserRoles.PROXY_ADMIN) + response = http_client.get("/get/latest_release_info") + assert response.status_code == 200 + assert response.json() == expected_body + assert client.urls == expected_urls + def test_repeated_requests_reuse_cache(self, http_client): - client = _RecordingClient([_github_response()]) + client = _minor_release() _override_dependencies(client, _fresh_cache(), LitellmUserRoles.PROXY_ADMIN) assert http_client.get("/get/latest_release_info").json() == EXPECTED_INFO assert http_client.get("/get/latest_release_info").json() == EXPECTED_INFO - assert len(client.calls) == 1 + assert len(client.calls) == 2 def test_rejects_unauthenticated_requests(self, monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.master_key", MASTER_KEY) 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 512f58380ea..57b68e6db05 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 @@ -1,6 +1,8 @@ import json import os +from collections.abc import Iterator from typing import Final +from unittest.mock import patch import pytest from fastapi.testclient import TestClient @@ -181,6 +183,12 @@ def mock_auth(): app.dependency_overrides.pop(user_api_key_auth, None) +@pytest.fixture(autouse=True) +def isolated_environment() -> Iterator[None]: + with patch.dict(os.environ): + yield + + @pytest.fixture(autouse=True) def fresh_settings_store(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server @@ -4345,3 +4353,105 @@ class TestMoyaiUrlSetting: await update_ui_settings({"moyai_url": None}, actor) assert persisted["moyai_url"] is None + + +@pytest.mark.parametrize("preset", ["today", "last_7_days", "last_30_days", "month_to_date", "year_to_date"]) +def test_usage_default_range_round_trips_without_overwriting_other_settings( + mock_auth, monkeypatch: pytest.MonkeyPatch, preset: str, +) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.config_resolvers import SettingsStore + + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=SimpleNamespace( + ui_settings={"disable_custom_api_keys": True}, + )) + prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server.proxy_config, "settings", SettingsStore("general_settings")) + actor: Final = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + app.dependency_overrides[user_api_key_auth] = lambda: actor + try: + initial: Final = client.get("/get/ui_settings") + assert initial.json()["values"]["usage_default_time_range"] == "last_7_days" + response: Final = client.patch("/update/ui_settings", json={"usage_default_time_range": preset}) + assert response.status_code == 200, response.text + stored: Final = json.loads(prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["update"]["ui_settings"]) + assert stored == {"disable_custom_api_keys": True, "usage_default_time_range": preset} + prisma.db.litellm_uisettings.find_unique.return_value = SimpleNamespace(ui_settings=stored) + fetched: Final = client.get("/get/ui_settings") + assert fetched.json()["values"]["usage_default_time_range"] == preset + assert fetched.json()["source"]["usage_default_time_range"] == "db" + finally: + app.dependency_overrides.clear() + + +@pytest.mark.parametrize("submitted", ["invalid", None, 7]) +def test_usage_default_range_rejects_invalid_values_without_writing( + mock_auth, monkeypatch: pytest.MonkeyPatch, submitted: object, +) -> None: + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + actor: Final = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + app.dependency_overrides[user_api_key_auth] = lambda: actor + try: + response: Final = client.patch("/update/ui_settings", json={"usage_default_time_range": submitted}) + assert response.status_code == 422, response.text + prisma.db.litellm_uisettings.upsert.assert_not_called() + finally: + app.dependency_overrides.clear() + + +@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) +def test_usage_default_range_cannot_be_changed_by_non_admins( + mock_auth, monkeypatch: pytest.MonkeyPatch, role: LitellmUserRoles, +) -> None: + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + actor: Final = UserAPIKeyAuth(user_id="test-user", user_role=role) + app.dependency_overrides[user_api_key_auth] = lambda: actor + try: + response: Final = client.patch("/update/ui_settings", json={"usage_default_time_range": "month_to_date"}) + assert response.status_code == 403, response.text + prisma.db.litellm_uisettings.upsert.assert_not_called() + finally: + app.dependency_overrides.clear() + + +def test_usage_default_range_reports_config_precedence(mock_auth, monkeypatch: pytest.MonkeyPatch) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy import proxy_server + from litellm.proxy.config_resolvers import SettingsStore + + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=SimpleNamespace( + ui_settings={"usage_default_time_range": "today"}, + )) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + settings: Final = SettingsStore("general_settings") + settings.load_yaml({"usage_default_time_range": "month_to_date"}) + monkeypatch.setattr(proxy_server.proxy_config, "settings", settings) + response: Final = client.get("/get/ui_settings") + assert response.status_code == 200, response.text + assert response.json()["values"]["usage_default_time_range"] == "month_to_date" + assert response.json()["source"]["usage_default_time_range"] == "config" diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py index 747e8da8043..e335b08f040 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py @@ -12,15 +12,118 @@ import asyncio import hashlib import json import logging +from collections.abc import Mapping +from contextlib import AbstractAsyncContextManager from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from litellm._service_logger import ServiceTypes +from litellm.models.user import LiteLLM_UserTable from litellm.proxy.db.log_db_metrics import record_db_io from litellm.proxy.utils import PrismaClient +from litellm.repositories.user_repository import free_sso_user_creation + + +class _SSOAdmissionUsers: + def __init__(self, count: int) -> None: + self.rows: Mapping[str, LiteLLM_UserTable] = { + f"user-{index}": LiteLLM_UserTable(user_id=f"user-{index}") + for index in range(count) + } + + async def find_unique( + self, where: Mapping[str, object] + ) -> LiteLLM_UserTable | None: + return self.rows.get(str(where["user_id"])) + + async def count(self, where: Mapping[str, object] | None = None) -> int: + if where is None: + return len(self.rows) + return sum((row.metadata or {}).get("scim_active") is False for row in self.rows.values()) + + async def upsert( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> LiteLLM_UserTable: + existing: Final = await self.find_unique(where) + if existing is not None: + return existing + created: Final = LiteLLM_UserTable.model_validate(data["create"]) + self.rows = {**self.rows, created.user_id: created} + return created + + +class _SSOAdmissionTransaction(AbstractAsyncContextManager["_SSOAdmissionTransaction"]): + def __init__(self, users: _SSOAdmissionUsers) -> None: + self.litellm_usertable: Final = users + + async def __aenter__(self) -> _SSOAdmissionTransaction: + return self + + async def __aexit__(self, *exc_info: object) -> None: + return None + + async def execute_raw(self, query: str) -> int: + return 0 + + async def query_raw(self, query: str) -> object: + return ({"locked": 1},) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("existing_count", "scope_id", "user_id", "enabled", "child_task", "refused"), + [ + (4, "new-user", "new-user", True, False, False), + (5, "new-user", "new-user", True, False, True), + (6, "user-0", "user-0", True, False, False), + (5, "sso-user", "new-user", True, False, False), + (5, "new-user", "new-user", True, True, False), + (5, None, "new-user", True, False, True), + (5, None, "new-user", False, False, False), + ], +) +async def test_user_insert_enforces_free_sso_admission_only_for_its_owner( + prisma_client: PrismaClient, + monkeypatch: pytest.MonkeyPatch, + existing_count: int, + scope_id: str | None, + user_id: str, + enabled: bool, + child_task: bool, + refused: bool, +) -> None: + users: Final = _SSOAdmissionUsers(existing_count) + prisma_client.db.litellm_usertable = users + prisma_client.db.tx.return_value = _SSOAdmissionTransaction(users) + monkeypatch.setattr("asyncio.sleep", AsyncMock(return_value=None)) + original: Final = users.rows.get(user_id) + + async def attempt() -> None: + if refused: + with pytest.raises(HTTPException) as error: + await prisma_client.insert_data(data={"user_id": user_id}, table_name="user") + assert error.value.status_code == 403 + assert user_id not in users.rows + elif child_task: + result: Final = await asyncio.create_task( + prisma_client.insert_data(data={"user_id": user_id}, table_name="user") + ) + assert result is users.rows[user_id] + else: + inserted: Final = await prisma_client.insert_data(data={"user_id": user_id}, table_name="user") + assert inserted is users.rows[user_id] + + with free_sso_user_creation(scope_id, enabled=enabled): + for _ in range(2): + await attempt() + + assert len(users.rows) == existing_count + int(not refused and original is None) + if original is not None: + assert users.rows[user_id] is original @pytest.mark.asyncio diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index e83efa21dcc..6557f0d1446 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -18,7 +18,7 @@ from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, ProxyLogging from litellm.types.utils import CachingDetails @@ -194,7 +194,9 @@ async def test_post_call_failure_hook_attributes_single_router_deployment( @pytest.mark.asyncio -async def test_pre_routing_reject_spend_log_keeps_public_model_group(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_pre_routing_reject_spend_log_keeps_public_model_group( + proxy_logging, make_user_api_key_auth, monkeypatch +): from litellm.proxy import proxy_server from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -1001,7 +1003,9 @@ async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: standard_logging_object = kwargs.get("standard_logging_object") - recorded.append(standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None) + recorded.append( + standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None + ) monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()]) logging_obj = LiteLLMLoggingObj( @@ -1029,3 +1033,52 @@ async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( await asyncio.sleep(0) await asyncio.sleep(0) assert recorded == ["failure"] + + +@pytest.mark.asyncio +async def test_handle_logging_proxy_only_path_flags_requests_that_died_before_request_setup( + proxy_logging, make_user_api_key_auth +): + """A reject before the pipeline set up logging (bad key, bad body) reaches the failure hook with no logging + object. The handler mints a ``litellm_call_id`` so error logs can carry it, and marks the request with a + server-owned sentinel so spend tracking can tell it apart from a guardrail block, which arrives with the + pipeline's logging object. ``request_data`` is the client body, so a client-sent call id or marker value must + not change either outcome.""" + rejected_before_setup: dict[str, object] = { + "model": "m", + "messages": [{"role": "user", "content": "x"}], + "metadata": {}, + "litellm_call_id": "client-supplied", + } + await proxy_logging._handle_logging_proxy_only_error( + request_data=rejected_before_setup, + user_api_key_dict=make_user_api_key_auth(), + route="/chat/completions", + original_exception=HTTPException(status_code=401, detail="bad key"), + ) + + logging_obj = MagicMock() + logging_obj.call_type = "acompletion" + logging_obj.model_call_details = {} + logging_obj.async_failure_handler = AsyncMock() + blocked_by_guardrail: dict[str, object] = { + "litellm_logging_obj": logging_obj, + "litellm_call_id": "call-1", + "model": "m", + "messages": [{"role": "user", "content": "x"}], + "metadata": {}, + REJECTED_BEFORE_REQUEST_SETUP_KEY: True, + } + await proxy_logging._handle_logging_proxy_only_error( + request_data=blocked_by_guardrail, + user_api_key_dict=make_user_api_key_auth(), + route="/chat/completions", + original_exception=HTTPException(status_code=400, detail="blocked"), + ) + + assert ( + rejected_before_setup.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP, + rejected_before_setup["litellm_call_id"], + blocked_by_guardrail.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP, + blocked_by_guardrail["litellm_call_id"], + ) == (True, "client-supplied", False, "call-1") diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index a08d44b9b4d..082a29b7001 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -6,6 +6,7 @@ import pytest from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, + SpendResetWrites, budget_cascade_unit_of_work, enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, @@ -39,6 +40,9 @@ class FakeBatch: self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls) self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) + def execute_raw(self, query: str, *args: object) -> None: + self.calls.append(("execute_raw", {"query": query}, {"args": args})) + async def commit(self) -> None: self.commit_count += 1 @@ -166,7 +170,7 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch( async with budget_cascade_unit_of_work(_new_batch) as uow: writes = [getattr(uow, field.name) for field in fields(uow)] for write in writes: - if isinstance(write, LinkedSpendResetWrites): + if isinstance(write, SpendResetWrites): write.queue_spend_zero(where={"budget_id": "budget-1"}) else: write.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) @@ -192,3 +196,18 @@ async def test_budget_cascade_raising_inside_block_skips_commit(): await _blow_up_mid_transaction() assert batch.commit_count == 0 + + +async def test_membership_total_spend_repair_is_queued_on_the_cascade_batch_ahead_of_the_reset(): + batch = FakeBatch() + + async with budget_cascade_unit_of_work(lambda: batch) as uow: + uow.team_memberships.queue_total_spend_repair(["budget-1", "budget-2"]) + uow.team_memberships.queue_spend_zero(where={"budget_id": "budget-1"}) + assert batch.commit_count == 0 + + assert batch.commit_count == 1 + repair, reset = batch.calls + assert repair[0] == "execute_raw" + assert repair[2] == {"args": (["budget-1", "budget-2"],)} + assert reset == ("litellm_teammembership.update_many", {"budget_id": "budget-1"}, {"spend": 0}) diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 94efef050a6..36df738c0bb 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -1,15 +1,17 @@ import json from collections.abc import Mapping, Sequence from itertools import chain -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from openai.types.responses.response_compaction_item import ResponseCompactionItem from pydantic import BaseModel import litellm +from litellm import CustomLLM from litellm.constants import REDACTED_BY_LITELLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( @@ -52,6 +54,7 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): instructions: str | None = None, request_kwargs: Mapping[str, object] | None = None, previous_response_client_ids: tuple[str, ...] = (), + input_replaces_history: bool = False, ) -> dict[str, object]: litellm_completion_request["litellm_trace_id"] = "session-trace" return litellm_completion_request @@ -561,3 +564,637 @@ def test_streamed_responses_over_anthropic_emit_the_responses_event_sequence(res assert deltas == "Argentina won." assert completed.response.status == "completed" assert completed.response.output[0].content[0].text == "Argentina won." + + +_ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages" +_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "The user pasted a long document and asked about it."} +_FIRST_INPUT: Final = ({"role": "user", "content": "What is the first word of the document?"},) +_NEXT_INPUT: Final = ({"role": "user", "content": "And the last word?"},) +_COMPACT_AT_50K: Final = ({"type": "compaction", "compact_threshold": 50000},) + + +def _anthropic_message(*content: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_compacted", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": list(content), + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + +def _replay(output: Sequence[BaseModel]) -> list[object]: + return [*_FIRST_INPUT, *(item.model_dump(exclude_none=True) for item in output), *_NEXT_INPUT] + + +def _sent_assistant_content(route: respx.Route) -> list[object]: + sent: Final = json.loads(route.calls.last.request.content) + return next(message["content"] for message in sent["messages"] if message["role"] == "assistant") + + +def test_responses_over_anthropic_return_the_compaction_block_and_send_it_back( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + first: Final = litellm.responses( + model=_MODEL, input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + litellm.responses( + model=_MODEL, input=_replay(first.output), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + + assert [item.type for item in first.output] == ["compaction", "message"] + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + + +def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_back( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_stream", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "compaction", "content": None}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "compaction_delta", "content": _COMPACTION_BLOCK["content"]}, + }, + {"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": "The."}}, + {"type": "content_block_stop", "index": 1}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) + item_events: Final = tuple( + event + for event in events + if event.type in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + ) + completed: Final = events[-1].response + litellm.responses( + model=_MODEL, input=_replay(completed.output), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + + assert [(event.type, event.output_index, event.item.type) for event in item_events] == [ + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 0, "compaction"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 1, "message"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 0, "compaction"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 1, "message"), + ] + assert item_events[0].item.encrypted_content == "" + assert item_events[0].item.id == item_events[2].item.id + assert [(item.type, getattr(item, "encrypted_content", None)) for item in completed.output[:1]] == [ + ("compaction", item_events[2].item.encrypted_content) + ] + assert completed.output[0].id == item_events[2].item.id + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + + +def test_streamed_responses_over_anthropic_answer_after_a_failed_compaction(respx_mock: respx.MockRouter) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_failed_compaction", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "compaction", "content": None}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "compaction_delta", "content": None}}, + {"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": "The."}}, + {"type": "content_block_stop", "index": 1}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) + + assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert [item.type for item in events[-1].response.output] == ["message"] + assert events[-1].response.output[0].content[0].text == "The." + + +@pytest.mark.asyncio +async def test_async_streamed_compaction_item_is_done_before_the_reasoning_that_follows_it() -> None: + anthropic: Final = httpx.MockTransport( + lambda _request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_thinking", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "compaction", "content": None}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "compaction_delta", "content": _COMPACTION_BLOCK["content"]}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "thinking", "thinking": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "thinking_delta", "thinking": "Recall."}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "signature_delta", "signature": "sig"}}, + {"type": "content_block_stop", "index": 1}, + {"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 2, "delta": {"type": "text_delta", "text": "The."}}, + {"type": "content_block_stop", "index": 2}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + stream: Final = await litellm.aresponses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + client=AsyncHTTPHandler(transport=anthropic), + ) + events: Final = [event async for event in stream] + done_items: Final = [ + (event.output_index, event.item.type) + for event in events + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE + ] + reasoning_done: Final = next( + index + for index, event in enumerate(events) + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and event.item.type == "reasoning" + ) + sequence_numbers: Final = [ + event.sequence_number for event in events[: reasoning_done + 1] if getattr(event, "sequence_number", None) + ] + + assert done_items == [(0, "compaction"), (1, "reasoning"), (2, "message")] + assert [item.type for item in events[-1].response.output] == ["compaction", "reasoning", "message"] + assert sequence_numbers == sorted(set(sequence_numbers)) + + +def _compaction_stream_events(*summaries: str) -> tuple[Mapping[str, object], ...]: + return tuple( + chain.from_iterable( + ( + { + "type": "content_block_start", + "index": 2 * index, + "content_block": {"type": "compaction", "content": None}, + }, + { + "type": "content_block_delta", + "index": 2 * index, + "delta": {"type": "compaction_delta", "content": summary}, + }, + {"type": "content_block_stop", "index": 2 * index}, + {"type": "content_block_start", "index": 2 * index + 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 2 * index + 1, "delta": {"type": "text_delta", "text": "Ok."}}, + {"type": "content_block_stop", "index": 2 * index + 1}, + ) + for index, summary in enumerate(summaries) + ) + ) + + +def test_streamed_compaction_item_done_carries_the_same_newest_block_as_the_completed_response( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_twice", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + *_compaction_stream_events("First summary.", "Second summary."), + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) + streamed: Final = tuple( + (event.item.id, json.loads(event.item.encrypted_content)["content"]) + for event in events + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and event.item.type == "compaction" + ) + completed: Final = tuple( + (item.id, json.loads(item.encrypted_content)["content"]) + for item in events[-1].response.output + if item.type == "compaction" + ) + + assert streamed == completed + assert [summary for _, summary in completed] == ["Second summary."] + + +def test_responses_over_anthropic_send_back_a_compaction_item_replayed_as_sdk_output( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + first: Final = litellm.responses( + model=_MODEL, input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + litellm.responses( + model=_MODEL, + input=[*_FIRST_INPUT, *first.output, *_NEXT_INPUT], + context_management=list(_COMPACT_AT_50K), + api_key="mock_api_key", + ) + + assert _sent_assistant_content(route)[0] == _COMPACTION_BLOCK + + +@pytest.mark.asyncio +async def test_router_with_encrypted_content_affinity_sends_the_compaction_block_back( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": _MODEL, "api_key": "mock_api_key"}, + "model_info": {"id": "deployment-a"}, + } + ], + optional_pre_call_checks=["encrypted_content_affinity"], + ) + + first: Final = await router.aresponses( + model="claude", input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K) + ) + await router.aresponses(model="claude", input=_replay(first.output), context_management=list(_COMPACT_AT_50K)) + + assert first.output[0].encrypted_content.startswith("litellm_enc:") + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + assert [key for key in json.loads(route.calls.last.request.content) if key.startswith("_")] == [] + + +_DOCUMENT_BEFORE_COMPACTION: Final = "The design document pasted before the compaction." +_TURN_BEFORE_COMPACTION: Final = { + "request_id": "chatcmpl-before-compaction", + "call_type": "aresponses", + "session_id": "session-compacted", + "proxy_server_request": { + "model": _MODEL, + "input": [ + {"role": "developer", "content": "Use metric units."}, + {"role": "user", "content": _DOCUMENT_BEFORE_COMPACTION}, + ], + }, + "response": { + "id": "chatcmpl-before-compaction", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Read it."}}], + }, +} + + +class _RecordingRawAnthropicMessages: + def __init__(self) -> None: + self.body: str = "" + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.body = request.content.decode() + return _anthropic_message({"type": "text", "text": "dog."}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +@pytest.mark.parametrize( + "compaction_item", + [ + pytest.param( + {"type": "compaction", "id": "cmp_caller", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + id="param_dict", + ), + pytest.param( + ResponseCompactionItem(type="compaction", id="cmp_caller", encrypted_content=json.dumps(_COMPACTION_BLOCK)), + id="sdk_output_item", + ), + ], +) +async def test_previous_response_id_with_a_compaction_item_in_input_replays_only_the_stored_instructions( + context_management: list[Mapping[str, object]] | None, + compaction_item: Mapping[str, object] | ResponseCompactionItem, +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_TURN_BEFORE_COMPACTION])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-before-compaction", + input=[compaction_item, *_NEXT_INPUT], + context_management=context_management, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + assert _DOCUMENT_BEFORE_COMPACTION not in anthropic.body + assert "Read it." not in anthropic.body + assert "Use metric units." in anthropic.body + assert _COMPACTION_BLOCK["content"] in anthropic.body + assert _NEXT_INPUT[0]["content"] in anthropic.body + + +class _RecordingCustomLLM(CustomLLM): + def __init__(self) -> None: + super().__init__() + self.sent_messages: list[object] = [] + + def completion(self, *args: object, **kwargs: object) -> ModelResponse: + self.sent_messages.extend(cast(list[object], kwargs["messages"])) + return litellm.completion( + model="gpt-5.6", messages=[{"role": "user", "content": "Hi."}], mock_response="dog." + ) + + +def test_responses_over_a_custom_provider_send_a_compaction_item_as_summary_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + custom_llm: Final = _RecordingCustomLLM() + monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) + monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "compaction-test-llm", "custom_handler": custom_llm}] + ) + compaction_item: Final = { + "type": "compaction", + "id": "cmp_earlier", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + + litellm.responses( + model="compaction-test-llm/my-model", + input=[*_FIRST_INPUT, {"role": "assistant", "content": "The."}, compaction_item, *_NEXT_INPUT], + context_management=list(_COMPACT_AT_50K), + drop_params=True, + ) + + assert custom_llm.sent_messages == [ + {"role": "user", "content": _COMPACTION_BLOCK["content"]}, + *_NEXT_INPUT, + ] + + +_WEATHER_CALL: Final = { + "id": "toolu_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, +} +_TOOL_TURN_BEFORE_COMPACTION: Final = { + **_TURN_BEFORE_COMPACTION, + "request_id": "chatcmpl-tool-turn", + "response": { + "id": "chatcmpl-tool-turn", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": [_WEATHER_CALL]}, + } + ], + }, +} + + +def _content_blocks(message: Mapping[str, object]) -> list[Mapping[str, object]]: + content: Final = message["content"] + return content if isinstance(content, list) else [{"type": "text", "text": content}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +async def test_previous_response_id_with_a_compaction_item_keeps_the_stored_call_its_tool_output_answers( + context_management: list[Mapping[str, object]] | None, +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + compaction_item: Final = { + "type": "compaction", + "id": "cmp_caller", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + tool_output: Final = {"type": "function_call_output", "call_id": "toolu_weather", "output": "18C and sunny"} + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_TOOL_TURN_BEFORE_COMPACTION])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-tool-turn", + input=[compaction_item, tool_output], + context_management=context_management, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + messages: Final = json.loads(anthropic.body)["messages"] + call_turn: Final = next( + index + for index, message in enumerate(messages) + if any(block["type"] == "tool_use" for block in _content_blocks(message)) + ) + tool_use: Final = next(block for block in _content_blocks(messages[call_turn]) if block["type"] == "tool_use") + assert (tool_use["id"], tool_use["input"]) == ("toolu_weather", {"city": "Paris"}) + assert _content_blocks(messages[call_turn + 1])[0]["type"] == "tool_result" + assert _content_blocks(messages[call_turn + 1])[0]["tool_use_id"] == "toolu_weather" + assert _COMPACTION_BLOCK["content"] in json.dumps(messages[: call_turn + 1]) + assert _DOCUMENT_BEFORE_COMPACTION not in anthropic.body + + +_OSAKA_CALL: Final = {**_WEATHER_CALL, "function": {"name": "get_weather", "arguments": '{"city": "Osaka"}'}} +_TOOL_TURN_REUSING_THE_CALL_ID: Final = { + **_TOOL_TURN_BEFORE_COMPACTION, + "request_id": "chatcmpl-osaka-turn", + "proxy_server_request": { + "model": _MODEL, + "input": [{"type": "function_call_output", "call_id": "toolu_weather", "output": "18C and sunny"}], + }, + "response": { + "id": "chatcmpl-osaka-turn", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": [_OSAKA_CALL]}, + } + ], + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +async def test_previous_response_id_with_a_compaction_item_restores_only_the_newest_stored_call_for_a_reused_call_id( + context_management: list[Mapping[str, object]] | None, +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + compaction_item: Final = { + "type": "compaction", + "id": "cmp_caller", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + tool_output: Final = {"type": "function_call_output", "call_id": "toolu_weather", "output": "12C and rain"} + spend_logs: Final = [_TOOL_TURN_BEFORE_COMPACTION, _TOOL_TURN_REUSING_THE_CALL_ID] + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient(spend_logs)): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-osaka-turn", + input=[compaction_item, tool_output], + context_management=context_management, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + messages: Final = json.loads(anthropic.body)["messages"] + blocks: Final = list( + chain.from_iterable( + ((index, block) for block in _content_blocks(message)) for index, message in enumerate(messages) + ) + ) + tool_uses: Final = [(index, block["id"], block["input"]) for index, block in blocks if block["type"] == "tool_use"] + tool_results: Final = [(index, block["tool_use_id"]) for index, block in blocks if block["type"] == "tool_result"] + assert [(tool_id, tool_input) for _, tool_id, tool_input in tool_uses] == [("toolu_weather", {"city": "Osaka"})] + assert tool_results == [(tool_uses[0][0] + 1, "toolu_weather")] + assert "12C and rain" in json.dumps(messages[tool_uses[0][0] + 1]) + + +def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_the_request_does_not_compact( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "dog."}) + ) + compaction_item: Final = { + "type": "compaction", + "id": "cmp_earlier", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + + litellm.responses( + model=_MODEL, + input=[*_FIRST_INPUT, {"role": "assistant", "content": "The."}, compaction_item, *_NEXT_INPUT], + api_key="mock_api_key", + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert sent["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": _COMPACTION_BLOCK["content"]}, + {"type": "text", "text": _NEXT_INPUT[0]["content"]}, + ], + } + ] diff --git a/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py b/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py new file mode 100644 index 00000000000..e0ef3fbb552 --- /dev/null +++ b/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py @@ -0,0 +1,87 @@ +import json +from typing import Final + +import pytest + +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_output_item, + holds_compaction_summary, + response_holds_compaction_block, + with_compaction_summary, +) +from litellm.responses.utils import ResponsesAPIRequestUtils + +_BLOCK: Final = {"type": "compaction", "content": "Turn one planned Tokyo."} +_ITEM: Final = {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps(_BLOCK)} + + +@pytest.mark.parametrize( + "newer_entry", + [ + pytest.param("not a block", id="not_a_mapping"), + pytest.param({"type": "compaction", "content": None}, id="summary_not_streamed_yet"), + pytest.param({"type": "compaction", "content": ""}, id="empty_summary"), + pytest.param({"type": "text", "content": "Hi."}, id="not_a_compaction"), + ], +) +def test_compaction_output_item_carries_the_newest_block_that_holds_a_summary(newer_entry: object) -> None: + item: Final = compaction_output_item({"compaction_blocks": [_BLOCK, newer_entry]}) + + assert item is not None + assert json.loads(item.encrypted_content) == _BLOCK + + +@pytest.mark.parametrize( + ("response", "holds"), + [ + pytest.param( + {"choices": [{"message": {"role": "assistant", "provider_specific_fields": {"compaction_blocks": [_BLOCK]}}}]}, + True, + id="chat_answer_compacted", + ), + pytest.param({"choices": [{"message": {"role": "assistant", "content": "Hi."}}]}, False, id="chat_answer"), + pytest.param({"object": "response", "output": [_ITEM]}, True, id="responses_output_compacted"), + pytest.param( + {"object": "response", "output": [{"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAA"}]}, + False, + id="responses_output_with_an_opaque_compaction", + ), + pytest.param({"choices": "not a list"}, False, id="malformed_choices"), + pytest.param(None, False, id="no_stored_response"), + ], +) +def test_response_holds_compaction_block_reads_chat_and_responses_shaped_rows(response: object, holds: bool) -> None: + assert response_holds_compaction_block(response) is holds + + +@pytest.mark.parametrize( + ("encrypted_content", "holds"), + [ + pytest.param(json.dumps(_BLOCK), True, id="litellm_block"), + pytest.param( + ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(json.dumps(_BLOCK), "deployment-1"), + True, + id="affinity_wrapped_block", + ), + pytest.param("gAAAAABopaque", False, id="openai_token"), + pytest.param(None, False, id="missing"), + ], +) +def test_only_a_compaction_item_litellm_can_decode_holds_a_summary(encrypted_content: object, holds: bool) -> None: + assert holds_compaction_summary([{"type": "compaction", "id": "cmp_1", "encrypted_content": encrypted_content}]) is ( + holds + ) + + +def test_with_compaction_summary_rewrites_the_summary_and_keeps_the_affinity_wrap() -> None: + wrapped: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(json.dumps(_BLOCK), "deployment-1") + + rewritten: Final = with_compaction_summary(wrapped, "[REDACTED]") + model_id, original = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(rewritten) + + assert model_id == "deployment-1" + assert json.loads(original) == {**_BLOCK, "content": "[REDACTED]"} + + +def test_with_compaction_summary_leaves_a_token_it_cannot_read_unchanged() -> None: + assert with_compaction_summary("gAAAAABopaque", "[REDACTED]") == "gAAAAABopaque" diff --git a/tests/unit/responses/litellm_completion_transformation/test_count_request.py b/tests/unit/responses/litellm_completion_transformation/test_count_request.py index 464bdb1a367..0477ae15b8e 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_count_request.py +++ b/tests/unit/responses/litellm_completion_transformation/test_count_request.py @@ -10,6 +10,7 @@ from litellm.responses.litellm_completion_transformation.count_request import ( ChatCountRequest, chat_count_request, ) +from litellm.responses.litellm_completion_transformation.transformation import ChatCompletionSession _MARK: Final[dict[str, JsonValue]] = {"type": "ephemeral"} _MARKED_RESPONSES_REQUEST: Final[dict[str, JsonValue]] = { @@ -125,3 +126,77 @@ async def test_claude_count_request_leaves_out_tools_the_deployment_drops() -> N assert isinstance(dropped, ChatCountRequest) assert dropped.tools is None assert dropped == await _claude_count({"input": [_QUESTION]}, policy) + + +_SUMMARY: Final = "Earlier turns covered the Tokyo trip." +_COMPACTED_INPUT: Final[list[JsonValue]] = [ + {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": _SUMMARY})}, + {"type": "message", "role": "user", "content": "And Osaka?"}, +] +_COMPACTION_REQUESTS: Final = ( + pytest.param({}, id="summary_route"), + pytest.param({"context_management": [{"type": "compaction", "compact_threshold": 50000}]}, id="native_route"), +) + + +async def _stored_turns_the_compaction_replaces(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession( + messages=[ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_0", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + ], + litellm_session_id=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("compaction_request", _COMPACTION_REQUESTS) +async def test_claude_count_request_counts_a_compaction_item_as_its_summary_text( + compaction_request: dict[str, JsonValue], +) -> None: + request: Final = await chat_count_request( + model="claude-opus-5-5", + custom_llm_provider="anthropic", + api_base=None, + responses_request={"model": "claude-opus-5-5", "input": _COMPACTED_INPUT, **compaction_request}, + policy=DropPolicy(drop_params=False, additional_drop_params=frozenset()), + ) + + assert isinstance(request, ChatCountRequest) + assert [(message.get("role"), message.get("content")) for message in request.messages] == [ + ("user", _SUMMARY), + ("user", "And Osaka?"), + ] + assert "compaction_blocks" not in json.dumps(request.messages, default=str) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("compaction_request", _COMPACTION_REQUESTS) +async def test_claude_count_request_leaves_out_the_stored_turns_a_compaction_item_replaces( + compaction_request: dict[str, JsonValue], +) -> None: + request: Final = await chat_count_request( + model="claude-opus-5-5", + custom_llm_provider="anthropic", + api_base=None, + responses_request={ + "model": "claude-opus-5-5", + "previous_response_id": "resp_before_the_compaction", + "input": _COMPACTED_INPUT, + **compaction_request, + }, + policy=DropPolicy(drop_params=False, additional_drop_params=frozenset()), + session_lookup=_stored_turns_the_compaction_replaces, + ) + + assert isinstance(request, ChatCountRequest) + counted: Final = json.dumps(request.messages, default=str) + assert _SUMMARY in counted and "And Osaka?" in counted + assert "Weather in SF?" not in counted + assert "call_0" not in counted diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index f34ca7d0ff8..023d0ea26fb 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, anthropic_messages_pt, ) +from litellm.responses.litellm_completion_transformation.reasoning_items import encode_thinking_blocks from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, ChatCompletionSession, @@ -3397,7 +3398,9 @@ class TestUsageTransformation: assert response_usage.input_tokens_details is not None assert response_usage.input_tokens_details.cached_tokens == 5 assert response_usage.input_tokens_details.text_tokens == 8 - assert "cache_write_tokens" not in response_usage.input_tokens_details.model_dump() + assert {"cache_write_tokens", "cache_write_tokens_details"}.isdisjoint( + response_usage.input_tokens_details.model_dump() + ) def test_transform_usage_with_cached_tokens_gemini(self): """Test that cached_tokens from Gemini are properly transformed to input_tokens_details""" @@ -6502,6 +6505,363 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No assert result["service_tier"] == "priority" +_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "Earlier turns covered the Tokyo trip."} + + +_COMPACTION_REQUESTED: Final = ResponsesAPIOptionalRequestParams( + context_management=[{"type": "compaction", "compact_threshold": 50000}] +) +_SUMMARY_AS_TEXT: Final = {"role": "user", "content": _COMPACTION_BLOCK["content"]} + + +def _anthropic_form_context_management(edit: Mapping[str, object]) -> ResponsesAPIOptionalRequestParams: + return cast( # cast-ok: Anthropic's native {"edits": [...]} form, which the param type does not declare + ResponsesAPIOptionalRequestParams, {"context_management": {"edits": [edit]}} + ) + + +@pytest.mark.parametrize( + ("model", "responses_api_request", "replayed_compaction"), + [ + ( + "anthropic/claude-sonnet-5-5", + _COMPACTION_REQUESTED, + { + "role": "assistant", + "content": None, + "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}, + }, + ), + ( + "anthropic/claude-sonnet-5-5", + _anthropic_form_context_management({"type": "compact_20260112"}), + { + "role": "assistant", + "content": None, + "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}, + }, + ), + ("anthropic/claude-sonnet-5-5", ResponsesAPIOptionalRequestParams(), _SUMMARY_AS_TEXT), + ( + "anthropic/claude-sonnet-5-5", + _anthropic_form_context_management({"type": "clear_tool_uses_20250919"}), + _SUMMARY_AS_TEXT, + ), + ("gemini/gemini-3.8-flash", _COMPACTION_REQUESTED, _SUMMARY_AS_TEXT), + ("bedrock/claude_platform/claude-sonnet-5-5", _COMPACTION_REQUESTED, _SUMMARY_AS_TEXT), + ], +) +def test_compaction_input_item_replays_as_a_compaction_block_only_when_anthropic_compacts( + model: str, responses_api_request: ResponsesAPIOptionalRequestParams, replayed_compaction: Mapping[str, object] +) -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAAABopaque"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "And Osaka?"}, + ], + ), + responses_api_request=responses_api_request, + ) + + assert [dict(message) for message in result["messages"]] == [ + replayed_compaction, + {"role": "user", "content": "And Osaka?"}, + ] + + +def _replayed_compacted_answer(model: str, context_management: object) -> list[tuple[object, object, object]]: + async def session_lookup(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession( + messages=[ + Message( + role="assistant", + content="Day one, Asakusa.", + provider_specific_fields={"compaction_blocks": [_COMPACTION_BLOCK]}, + ), + {"role": "user", "content": "Day two?"}, + Message(role="assistant", content="Ueno."), + ], + litellm_session_id=None, + ) + + request: Final = asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_compacted", + litellm_completion_request={ + "model": model, + "messages": [{"role": "user", "content": "And Osaka?"}], + **({"context_management": context_management} if context_management else {}), + }, + instructions=None, + session_lookup=session_lookup, + ) + ) + return [ + ( + message.get("role"), + message.get("content"), + (message.get("provider_specific_fields") or {}).get("compaction_blocks"), + ) + for message in request["messages"] + ] + + +_TURNS_AFTER_THE_COMPACTED_ANSWER: Final = [ + ("user", "Day two?", None), + ("assistant", "Ueno.", None), + ("user", "And Osaka?", None), +] +_ANSWER_AFTER_SUMMARY: Final = [ + ("user", _COMPACTION_BLOCK["content"], None), + ("assistant", "Day one, Asakusa.", None), + *_TURNS_AFTER_THE_COMPACTED_ANSWER, +] + + +@pytest.mark.parametrize( + ("model", "context_management", "replayed"), + [ + ( + "anthropic/claude-sonnet-5-5", + _COMPACTION_REQUESTED["context_management"], + [("assistant", "Day one, Asakusa.", [_COMPACTION_BLOCK]), *_TURNS_AFTER_THE_COMPACTED_ANSWER], + ), + ("anthropic/claude-sonnet-5-5", None, _ANSWER_AFTER_SUMMARY), + ( + "bedrock/claude_platform/claude-sonnet-5-5", + _COMPACTION_REQUESTED["context_management"], + _ANSWER_AFTER_SUMMARY, + ), + ("gpt-5.6", _COMPACTION_REQUESTED["context_management"], _ANSWER_AFTER_SUMMARY), + ], +) +def test_stored_compacted_answer_replays_its_block_only_when_anthropic_compacts( + model: str, context_management: object, replayed: list[tuple[object, object, object]] +) -> None: + assert _replayed_compacted_answer(model, context_management) == replayed + + +def test_summary_replay_drops_the_history_its_compaction_item_replaces() -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="gemini/gemini-3.8-flash", + input=cast( + ResponseInputParam, + [ + {"role": "system", "content": "Answer briefly."}, + {"role": "user", "content": "Plan a Tokyo trip."}, + {"role": "assistant", "content": "Day one, Asakusa."}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "And Osaka?"}, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + ) + + assert [dict(message) for message in result["messages"]] == [ + {"role": "system", "content": "Answer briefly."}, + _SUMMARY_AS_TEXT, + {"role": "user", "content": "And Osaka?"}, + ] + + + +@pytest.mark.parametrize( + ("call", "output"), + [ + pytest.param( + {"type": "function_call", "call_id": "call_weather", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_weather", "output": "Rain."}, + id="function_call", + ), + pytest.param( + {"type": "custom_tool_call", "call_id": "call_patch", "name": "apply_patch", "input": "*** Begin Patch"}, + {"type": "custom_tool_call_output", "call_id": "call_patch", "output": "Done."}, + id="custom_tool_call", + ), + ], +) +def test_summary_replay_keeps_a_call_its_compaction_item_replaces_when_a_later_output_answers_it( + call: Mapping[str, str], output: Mapping[str, str] +) -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="gemini/gemini-3.8-flash", + input=cast( + ResponseInputParam, + [ + {"role": "system", "content": "Answer briefly."}, + {"role": "user", "content": "Plan a Tokyo trip."}, + call, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + output, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + )["messages"] + + assert [message["role"] for message in messages] == ["system", "user", "assistant", "tool"] + assert dict(messages[1]) == _SUMMARY_AS_TEXT + assert [tool_call["id"] for tool_call in messages[2]["tool_calls"]] == [output["call_id"]] + assert messages[3]["tool_call_id"] == output["call_id"] + + +_ANSWERED_THINKING: Final = {"type": "thinking", "thinking": "Need Osaka weather.", "signature": "sig_osaka"} +_CHECKING_OSAKA: Final = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Checking Osaka."}], +} + + +def _reasoning_item(item_id: str, thinking: Mapping[str, str]) -> Mapping[str, object]: + return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": encode_thinking_blocks([thinking])} + + +def test_summary_replay_keeps_only_the_reasoning_of_the_turn_whose_call_a_later_output_answers() -> None: + tokyo_thinking: Final = {"type": "thinking", "thinking": "Need Tokyo weather.", "signature": "sig_tokyo"} + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + _reasoning_item("rs_tokyo", tokyo_thinking), + {"type": "function_call", "call_id": "call_tokyo", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_tokyo", "output": "Sun."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_osaka", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call", "call_id": "call_kyoto", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_kyoto", "output": "Clouds."}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_osaka", "output": "Rain."}, + ], + ), + responses_api_request=ResponsesAPIOptionalRequestParams(), + )["messages"] + + assert [message["role"] for message in messages] == ["user", "assistant", "tool"] + assert dict(messages[0]) == _SUMMARY_AS_TEXT + assert messages[1]["thinking_blocks"] == [_ANSWERED_THINKING] + assert [tool_call["id"] for tool_call in messages[1]["tool_calls"]] == ["call_osaka"] + + +def test_native_compaction_replay_moves_a_call_its_compaction_item_splits_from_its_output_after_the_block() -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_osaka", "name": "get_weather", "arguments": "{}"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_osaka", "output": "Rain."}, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + )["messages"] + + assert [message["role"] for message in messages] == ["user", "assistant", "assistant", "tool"] + assert messages[0]["content"] == "Plan a Tokyo trip." + assert messages[1]["provider_specific_fields"] == {"compaction_blocks": [_COMPACTION_BLOCK]} + assert messages[2]["thinking_blocks"] == [_ANSWERED_THINKING] + assert [tool_call["id"] for tool_call in messages[2]["tool_calls"]] == ["call_osaka"] + assert messages[3]["tool_call_id"] == "call_osaka" + + +@pytest.mark.parametrize( + "responses_api_request", + [ + pytest.param(_COMPACTION_REQUESTED, id="native_compaction_block"), + pytest.param(ResponsesAPIOptionalRequestParams(), id="summary_as_text"), + ], +) +def test_compaction_replay_invents_no_call_for_an_output_whose_call_was_never_stored( + responses_api_request: ResponsesAPIOptionalRequestParams, +) -> None: + async def no_stored_turns(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession(messages=[], litellm_session_id=None) + + chat_request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_never_stored", "output": "Rain."}, + ], + ), + responses_api_request=ResponsesAPIOptionalRequestParams( + **responses_api_request, + tools=[{"type": "function", "name": "get_weather", "parameters": {"type": "object"}, "strict": False}], + ), + ) + + request: Final = asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_unknown", + litellm_completion_request=chat_request, + instructions=None, + session_lookup=no_stored_turns, + ) + ) + + assert not any(message.get("tool_calls") for message in request["messages"]) + + +@pytest.mark.parametrize( + ("responses_api_request", "expected_turns"), + [ + pytest.param( + _COMPACTION_REQUESTED, + [("user", ()), ("assistant", ("call_0",)), ("tool", "Sunny."), ("assistant", ()), ("assistant", ("call_0",)), ("tool", "Rain.")], + id="native_compaction_block", + ), + pytest.param( + ResponsesAPIOptionalRequestParams(), + [("user", ()), ("assistant", ("call_0",)), ("tool", "Rain.")], + id="summary_as_text", + ), + ], +) +def test_compaction_replay_moves_only_the_call_a_later_output_answers_when_an_earlier_call_reused_its_id( + responses_api_request: ResponsesAPIOptionalRequestParams, expected_turns: list[tuple[str, object]] +) -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + {"type": "function_call", "call_id": "call_0", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_0", "output": "Sunny."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_0", "name": "get_weather", "arguments": "{}"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_0", "output": "Rain."}, + ], + ), + responses_api_request=responses_api_request, + )["messages"] + + assert [ + ( + message["role"], + message["content"] + if message["role"] == "tool" + else tuple(tool_call["id"] for tool_call in message.get("tool_calls") or ()), + ) + for message in messages + ] == expected_turns + assert messages[-2]["thinking_blocks"] == [_ANSWERED_THINKING] + + def _continued_chain_request(request_kwargs: Mapping[str, object]) -> Mapping[str, object]: async def session_lookup(previous_response_id: str) -> ChatCompletionSession: assert previous_response_id == "resp_turn_1" @@ -6603,3 +6963,47 @@ async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chun assert incomplete_event.response.usage is not None assert incomplete_event.response.usage.output_tokens == litellm.token_counter(model="gpt-4o-mini", text=relayed_text) assert incomplete_event.response.usage.input_tokens > 0 + + +async def _input_tokens_billed_on_disconnect(request_input: ResponseInputParam) -> int: + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + chat_stream: Final = await litellm.acompletion( + model="anthropic/claude-sonnet-5-5", + messages=[{"role": "user", "content": "Which day is driest?"}], + mock_response="Thursday looks driest, so plan the walking tour then.", + stream=True, + api_key="test-key", + ) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="anthropic/claude-sonnet-5-5", + litellm_custom_stream_wrapper=chat_stream, + request_input=request_input, + responses_api_request={}, + custom_llm_provider="anthropic", + ) + _ = [await iterator.__anext__() for _ in range(6)] + incomplete_event: Final = iterator.partial_response_on_disconnect() + assert incomplete_event is not None and incomplete_event.response.usage is not None + return incomplete_event.response.usage.input_tokens + + +@pytest.mark.asyncio +async def test_bridge_partial_response_on_disconnect_bills_only_the_input_since_the_newest_compaction_item(): + since_compaction: Final = cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "Which day is driest?"}, + ], + ) + history_it_replaces: Final = cast( + ResponseInputParam, + [{"role": "user", "content": "Plan a Tokyo trip, day by day, with every museum and food stall. " * 200}], + ) + + assert await _input_tokens_billed_on_disconnect( + cast(ResponseInputParam, [*history_it_replaces, *since_compaction]) + ) == await _input_tokens_billed_on_disconnect(since_compaction) diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 7a146c1b665..0eb1aa26a65 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,6 +1,7 @@ import json import sqlite3 from collections.abc import Sequence +from itertools import chain from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, patch @@ -18,6 +19,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, _normalize_redacted_tool_call_arguments, ) +from litellm.responses.litellm_completion_transformation.transformation import ChatCompletionSession from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse @@ -233,11 +235,10 @@ async def test_e2e_cold_storage_successful_retrieval(): patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, patch("litellm.cold_storage_custom_logger", return_value="s3"), ): - # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs - mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = ( - AsyncMock(return_value=full_proxy_request) + mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = AsyncMock( + return_value=full_proxy_request ) # Call the main function @@ -291,7 +292,6 @@ async def test_e2e_cold_storage_fallback_to_truncated_payload(): ) as mock_get_spend_logs, patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, ): - # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs @@ -358,41 +358,25 @@ async def test_should_check_cold_storage_for_full_payload(): with patch("litellm.cold_storage_custom_logger", return_value="s3"): # Test case 1: Should return True for truncated content - result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_with_truncated_pdf - ) - assert ( - result1 == True - ), "Should return True for proxy request with truncated PDF content" + result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result1 == True, "Should return True for proxy request with truncated PDF content" # Test case 2: Should return False for regular content - result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_regular - ) - assert ( - result2 == False - ), "Should return False for regular proxy request without truncation" + result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_regular) + assert result2 == False, "Should return False for regular proxy request without truncation" # Test case 3: Should return True for empty request - result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_empty - ) + result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_empty) assert result3 == True, "Should return True for empty proxy request" # Test case 4: Should return True for None request - result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_none - ) + result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_none) assert result4 == True, "Should return True for None proxy request" # Test case 5: Should return False when cold storage is not configured with patch.object(litellm, "cold_storage_custom_logger", None): - result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_with_truncated_pdf - ) - assert ( - result5 == False - ), "Should return False when cold storage is not configured, even with truncated content" + result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result5 == False, "Should return False when cold storage is not configured, even with truncated content" @pytest.mark.asyncio @@ -422,9 +406,7 @@ async def test_get_chat_completion_message_history_empty_response_dict(): } ] - with patch.object( - ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id" - ) as mock_get_spend_logs: + with patch.object(ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id") as mock_get_spend_logs: mock_get_spend_logs.return_value = mock_spend_logs # Call the function @@ -518,9 +500,7 @@ async def test_message_history_reconstructs_list_shaped_input(): ) as mock_get_spend_logs: mock_get_spend_logs.return_value = mock_spend_logs - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id - ) + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(request_id) messages = result["messages"] assert [(message.get("role"), message.get("content")) for message in messages] == [ @@ -549,9 +529,7 @@ async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flus fake_prisma_client = _FakePrismaClient(results=[[], [spend_log]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id - ) + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(request_id) messages = result["messages"] assert [(message.get("role"), message.get("content")) for message in messages] == [ @@ -618,14 +596,10 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", response_id=request_id, ) - fake_prisma_client = _FakePrismaClient( - results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]] - ) + fake_prisma_client = _FakePrismaClient(results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( - encoded_response_id - ) + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(encoded_response_id) assert len(fake_prisma_client.db.calls) == 1 assert request_id in fake_prisma_client.db.calls[0] @@ -724,8 +698,9 @@ async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( """ fake_prisma_client = _FakePrismaClient(results=[]) - with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), patch( - "litellm.proxy.proxy_server.disable_spend_logs", True + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), + patch("litellm.proxy.proxy_server.disable_spend_logs", True), ): spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( "chatcmpl-does-not-exist" @@ -741,7 +716,11 @@ def test_normalize_redacted_arguments_skips_custom_tool_calls(): content=None, tool_calls=[ {"id": "call_c", "type": "custom", "custom": {"name": "run_code", "input": "print(1)"}}, - {"id": "call_f", "type": "function", "function": {"name": "get_weather", "arguments": "redacted-by-litellm"}}, + { + "id": "call_f", + "type": "function", + "function": {"name": "get_weather", "arguments": "redacted-by-litellm"}, + }, ], ) @@ -796,11 +775,13 @@ async def test_message_history_normalizes_redacted_tool_call_arguments(): } ] - with patch.object( # test-quality-ok: the handler has no DI seam for the spend-log fetch; every test in this file stubs this same boundary - ResponsesSessionHandler, - "get_all_spend_logs_for_previous_response_id", - new_callable=AsyncMock, - ) as mock_get_spend_logs: + with ( + patch.object( # test-quality-ok: the handler has no DI seam for the spend-log fetch; every test in this file stubs this same boundary + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + ) as mock_get_spend_logs + ): mock_get_spend_logs.return_value = mock_spend_logs result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( @@ -878,6 +859,286 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: assert json.loads(tool_message["content"]) == function_output +def _turn_spend_log( + turn_input: str | list[dict[str, str]], answer: str, compaction_blocks: list[dict[str, str]] | None = None +) -> dict[str, object]: + compaction_fields: Final = ( + {"provider_specific_fields": {"compaction_blocks": compaction_blocks}} if compaction_blocks else {} + ) + return { + "request_id": f"resp_{len(turn_input)}", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": turn_input}), + "response": { + "id": f"chatcmpl-{len(turn_input)}", + "object": "chat.completion", + "choices": [{"index": 0, "message": {"role": "assistant", "content": answer, **compaction_fields}}], + }, + } + + +def _role_and_content(message: object) -> tuple[object, object]: + if isinstance(message, dict): + return message.get("role"), message.get("content") + return getattr(message, "role", None), getattr(message, "content", None) + + +async def _history_from_stored_turns( + spend_logs: tuple[dict[str, object], ...], previous_response_id: str +) -> ChatCompletionSession: + with patch("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=_StoredSpendLogs(spend_logs))): + return await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + previous_response_id + ) + + +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_the_newest_turn_that_sent_a_compaction_item(): + first_turn: Final = [{"role": "user", "content": "Plan a Tokyo trip."}] + compaction_item: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), + } + second_turn: Final = [ + *first_turn, + {"role": "assistant", "content": "Day one, Asakusa."}, + compaction_item, + {"role": "user", "content": "And Osaka?"}, + ] + spend_logs: Final = ( + _turn_spend_log(first_turn, "Day one, Asakusa."), + _turn_spend_log(second_turn, "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_4") + + assert [_role_and_content(message) for message in session["messages"]] == [ + ("user", "Turn one planned Tokyo."), + ("user", "And Osaka?"), + ("assistant", "Dotonbori."), + ] + + +@pytest.mark.parametrize( + "first_turn_input", + [pytest.param([{"role": "user", "content": "Hi."}], id="item_list"), pytest.param("Hi.", id="plain_string")], +) +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_the_newest_answer_that_compacted( + first_turn_input: str | list[dict[str, str]], +): + compaction_blocks: Final = [{"type": "compaction", "content": "Turn one planned Tokyo."}] + spend_logs: Final = ( + _turn_spend_log(first_turn_input, "Hello."), + _turn_spend_log( + [{"role": "developer", "content": "Answer briefly."}, {"role": "user", "content": "Plan a Tokyo trip."}], + "Day one, Asakusa.", + compaction_blocks, + ), + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + assert [_role_and_content(message) for message in session["messages"]] == [ + ("developer", "Answer briefly."), + ("assistant", "Day one, Asakusa."), + ("user", "And Osaka?"), + ("assistant", "Dotonbori."), + ] + assert getattr(session["messages"][1], "provider_specific_fields", None) == {"compaction_blocks": compaction_blocks} + + +_COMPACTION_ITEM_SENT: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} + + +@pytest.mark.parametrize( + "compacting_turn", + [ + pytest.param( + _turn_spend_log( + [{"role": "user", "content": "Plan a Tokyo trip."}], + "Day one, Asakusa.", + [{"type": "compaction", "content": "Turn one planned Tokyo."}], + ), + id="answer_compacted", + ), + pytest.param( + _turn_spend_log( + [_COMPACTION_ITEM_SENT, {"role": "user", "content": "Plan a Tokyo trip."}], "Day one, Asakusa." + ), + id="input_sent_a_compaction_item", + ), + ], +) +@pytest.mark.asyncio +async def test_previous_response_history_keeps_instructions_from_turns_the_compaction_covers( + compacting_turn: dict[str, object], +): + spend_logs: Final = ( + _turn_spend_log( + [{"role": "developer", "content": "Use metric units."}, {"role": "user", "content": "Hi."}], "Hello." + ), + compacting_turn, + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + roles_and_contents: Final = [_role_and_content(message) for message in session["messages"]] + assert roles_and_contents[0] == ("developer", "Use metric units.") + assert ("user", "Hi.") not in roles_and_contents + assert ("assistant", "Hello.") not in roles_and_contents + assert roles_and_contents[-2:] == [("user", "And Osaka?"), ("assistant", "Dotonbori.")] + + +@pytest.mark.asyncio +async def test_previous_response_history_drops_a_tool_output_carrying_a_system_role_the_compaction_covers(): + spend_logs: Final = ( + _turn_spend_log( + [ + {"role": "user", "content": "Look up the code."}, + {"type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "4242", "role": "system"}, + ], + "Found it.", + ), + _turn_spend_log([_COMPACTION_ITEM_SENT, {"role": "user", "content": "Plan a Tokyo trip."}], "Day one."), + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + contents: Final = [str(_role_and_content(message)[1]) for message in session["messages"]] + assert not any("4242" in content for content in contents) + assert contents[-2:] == ["And Osaka?", "Dotonbori."] + + +def _output_message(text: str) -> dict[str, object]: + return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} + + +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_a_compaction_item_in_a_stored_responses_output(): + responses_shaped_turn: Final = { + "request_id": "resp_mcp", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": [{"role": "user", "content": "Plan a Tokyo trip."}]}), + "response": { + "id": "resp_mcp", + "object": "response", + "output": [_output_message("Checking the map."), _COMPACTION_ITEM_SENT, _output_message("Day one.")], + }, + } + spend_logs: Final = ( + _turn_spend_log([{"role": "user", "content": "Hi."}], "Hello."), + responses_shaped_turn, + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + replayed: Final = json.dumps([_role_and_content(message) for message in session["messages"]], default=str) + assert "Hello." not in replayed + assert "Plan a Tokyo trip." not in replayed + assert "Checking the map." not in replayed + assert "Turn one planned Tokyo." in replayed + assert [str(_role_and_content(message)[1]) for message in session["messages"]][-2:] == [ + "And Osaka?", + "Dotonbori.", + ] + + +def _tool_call_ids(message: object) -> tuple[str, ...]: + tool_calls: Final = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) + return tuple(call["id"] if isinstance(call, dict) else call.id for call in tool_calls or ()) + + +_WEATHER_CALL_TURN: Final = { + "request_id": "resp_weather", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": [{"role": "user", "content": "Weather in Paris?"}]}), + "response": { + "id": "chatcmpl-weather", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + } + ], + }, +} + + +def _compacted_weather_answer(request_id: str) -> dict[str, object]: + return { + **_turn_spend_log( + [_COMPACTION_ITEM_SENT, {"type": "function_call_output", "call_id": "call_1", "output": "sunny"}], + "It is sunny.", + ), + "request_id": request_id, + } + + +@pytest.mark.asyncio +async def test_previous_response_history_ignores_a_compaction_item_sent_after_the_referenced_response(): + session: Final = await _history_from_stored_turns( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send")), "resp_weather" + ) + + assert _role_and_content(session["messages"][0]) == ("user", "Weather in Paris?") + assert tuple(chain.from_iterable(map(_tool_call_ids, session["messages"]))) == ("call_1",) + + +@pytest.mark.asyncio +async def test_a_turn_after_a_compaction_item_keeps_the_stored_call_its_tool_output_answers(): + session: Final = await _history_from_stored_turns( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_compacted")), "resp_compacted" + ) + + roles: Final = [_role_and_content(message)[0] for message in session["messages"]] + assert roles == ["user", "assistant", "tool", "assistant"] + assert _tool_call_ids(session["messages"][1]) == ("call_1",) + assert ("user", "Weather in Paris?") not in map(_role_and_content, session["messages"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend_logs", "previous_response_id", "replayed_roles"), + [ + ((_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send")), "resp_weather", ["user", "assistant"]), + ( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send"), _compacted_weather_answer("resp_retry")), + "resp_first_send", + ["user", "assistant", "tool", "assistant"], + ), + ], +) +async def test_a_compacted_session_replays_no_turn_sent_after_the_referenced_response( + spend_logs: tuple[dict[str, object], ...], previous_response_id: str, replayed_roles: list[str] +): + session: Final = await _history_from_stored_turns(spend_logs, previous_response_id) + + assert [_role_and_content(message)[0] for message in session["messages"]] == replayed_roles + + class _SqliteBackedPrismaDB: def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: self._connection = sqlite3.connect(":memory:") @@ -889,7 +1150,9 @@ class _SqliteBackedPrismaDB: self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: - sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace("$1", "?") + sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace( + "$1", "?" + ) sqlite_args: Final = tuple(json.dumps(arg) if isinstance(arg, list) else arg for arg in args) return [dict(row) for row in self._connection.execute(sqlite_query, sqlite_args)] diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 1a3430224d1..ed90acbd2fa 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1524,3 +1524,37 @@ async def test_caller_function_named_like_the_search_tool_streams_as_function_ca assert [(item["type"], item["name"]) for item in output] == [("function_call", "tool_search_tool_regex")] assert "response.function_call_arguments.done" in [event["type"] for event in events] + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_a_compaction_block_without_its_summary_streams_no_compaction_item(sync_mode: bool): + deltas: Final = ( + Delta( + role="assistant", + content="", + provider_specific_fields={"compaction_blocks": [{"type": "compaction", "content": None}]}, + ), + Delta(content="Ok."), + ) + chunks: Final = tuple( + ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-sonnet-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices(index=0, delta=delta, finish_reason="stop" if index == len(deltas) - 1 else None) + ], + ) + for index, delta in enumerate(deltas) + ) + iterator: Final = _build_iterator(chunks) + events: Final = ( + [event.model_dump(exclude_none=True) for event in iterator] + if sync_mode + else [event.model_dump(exclude_none=True) async for event in iterator] + ) + + assert [event["item"]["type"] for event in events if "item" in event] == ["message", "message"] + assert [item["type"] for item in events[-1]["response"]["output"]] == ["message"] diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 60d4d2e89b5..285d90ef8de 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1,7 +1,9 @@ import asyncio +import gc import json import logging import sys +import warnings from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager, nullcontext from datetime import datetime @@ -18,20 +20,23 @@ from mcp.types import Tool as MCPTool # Import required modules import litellm +from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.langfuse.langfuse import langfuse_trace_id, langfuse_trace_metadata from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.xai.common_utils import xai_reported_cost_in_usd from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk -from litellm.proxy.utils import _failure_fields_to_lift +from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_fields_to_lift from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.responses.mcp.mcp_streaming_iterator import MCPEnhancedStreamingIterator from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamingResponse, OpenAIMcpServerTool, ToolParam, ) -from litellm.types.utils import StandardLoggingPayload, Usage +from litellm.types.utils import ModelResponse, StandardLoggingPayload, Usage class MockUserAPIKeyAuth: @@ -1518,25 +1523,39 @@ def _openai_answer_round() -> bytes: ) -def _anthropic_round(message_id: str, block: Mapping[str, object], delta: Mapping[str, object], stop: str) -> bytes: +def _anthropic_message_start( + message_id: str, model: str, input_tokens: int, cache_usage: Mapping[str, object] = MappingProxyType({}) +) -> dict[str, object]: + return { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": input_tokens, "output_tokens": 1, **cache_usage}, + }, + } + + +def _anthropic_round( + message_id: str, + block: Mapping[str, object], + delta: Mapping[str, object], + stop: str, + model: str = "claude-haiku-4-5", + usage: tuple[int, int] = (5, 5), + cache_usage: Mapping[str, object] = MappingProxyType({}), +) -> bytes: return _sse( ( - { - "type": "message_start", - "message": { - "id": message_id, - "type": "message", - "role": "assistant", - "model": "claude-haiku-4-5", - "content": [], - "stop_reason": None, - "usage": {"input_tokens": 5, "output_tokens": 1}, - }, - }, + _anthropic_message_start(message_id, model, usage[0], cache_usage), {"type": "content_block_start", "index": 0, "content_block": dict(block)}, {"type": "content_block_delta", "index": 0, "delta": dict(delta)}, {"type": "content_block_stop", "index": 0}, - {"type": "message_delta", "delta": {"stop_reason": stop}, "usage": {"output_tokens": 5}}, + {"type": "message_delta", "delta": {"stop_reason": stop}, "usage": {"output_tokens": usage[1]}}, {"type": "message_stop"}, ) ) @@ -1586,6 +1605,7 @@ class _RoundRecorder(CustomLogger): def __init__(self) -> None: super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () + self.results: tuple[object, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () self.hooked: tuple[object, ...] = () @@ -1593,6 +1613,7 @@ class _RoundRecorder(CustomLogger): self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + self.results = (*self.results, response_obj) async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime @@ -1624,12 +1645,15 @@ _ECHO_MCP_TOOL: Final = MappingProxyType( def _serve_responses( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, responses: Sequence[httpx.Response] + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + responses: Sequence[httpx.Response], + url_regex: str = r"https://mcp-rounds\.example\.com/.*responses", ) -> respx.Route: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() - return respx_mock.post(url__regex=r"https://mcp-rounds\.example\.com/.*responses").mock(side_effect=list(responses)) + return respx_mock.post(url__regex=url_regex).mock(side_effect=list(responses)) def _serve_rounds( @@ -1957,7 +1981,22 @@ async def test_non_stream_auto_execute_failing_every_retry_bills_each_paid_first failure_row_fields: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) billed_usage: Final = failure_row_fields["combined_usage_object"] assert isinstance(billed_usage, Usage) - assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == billed_tokens + estimated_failure_usage: Final = ( + _estimate_dispatched_failure_usage( + model="gpt-5.6", + request_input=logging_obj.model_call_details.get("messages"), + system_input=None, + ) + if not retried_first_rounds + else None + ) + expected_prompt_tokens: Final = billed_tokens[0] + ( + estimated_failure_usage.prompt_tokens if estimated_failure_usage is not None else 0 + ) + assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == ( + expected_prompt_tokens, + billed_tokens[1], + ) assert failure_row_fields["response_cost"] == pytest.approx( _round_cost(_ROUND_ONE_BODY) + sum(_round_cost(body) for body in retried_first_rounds) ) @@ -2346,8 +2385,11 @@ def _first_index(events: Sequence[object], event_type: str) -> int: async def _scripted_mcp_stream( - model: str, upstream_url: str, rounds: Sequence[bytes] -) -> tuple[tuple[object, ...], respx.Route, AsyncMock, AsyncMock]: + model: str, + upstream_url: str, + rounds: Sequence[bytes], + max_agentic_loops: int | None = None, +) -> tuple[tuple[object, ...], respx.Route, AsyncMock, AsyncMock, MCPEnhancedStreamingIterator]: upstream: Final = respx.post(upstream_url).mock( side_effect=[httpx.Response(200, content=body, headers=SSE_HEADERS) for body in rounds] ) @@ -2364,9 +2406,11 @@ async def _scripted_mcp_stream( tool_choice="required", input=[{"role": "user", "type": "message", "content": "What's the weather in San Francisco?"}], stream=True, + **({"max_agentic_loops": max_agentic_loops} if max_agentic_loops is not None else {}), ) + assert isinstance(response, MCPEnhancedStreamingIterator) events: Final = tuple([event async for event in response]) - return events, upstream, fetch, execute + return events, upstream, fetch, execute, response @pytest.mark.asyncio @@ -2379,7 +2423,7 @@ async def test_streaming_responses_api_with_mcp_tools( model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS[provider] with caplog.at_level(logging.ERROR): - events, upstream, fetch, execute = await _scripted_mcp_stream( + events, upstream, fetch, execute, _ = await _scripted_mcp_stream( model, upstream_url, (tool_round(), answer_round()) ) @@ -2398,6 +2442,22 @@ async def test_streaming_responses_api_with_mcp_tools( assert not [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] +@pytest.mark.asyncio +@respx.mock +async def test_streaming_mcp_iterator_uses_max_agentic_loops_from_request(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS["openai"] + + _, _, _, _, iterator = await _scripted_mcp_stream( + model, + upstream_url, + (tool_round(), answer_round()), + max_agentic_loops=7, + ) + + assert iterator.max_tool_call_rounds == 7 + + @pytest.mark.asyncio @respx.mock async def test_streaming_mcp_event_order_and_response_id_consistency( @@ -2407,7 +2467,7 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS["openai"] with caplog.at_level(logging.ERROR): - events, _, _, _ = await _scripted_mcp_stream(model, upstream_url, (tool_round(), answer_round())) + events, _, _, _, _ = await _scripted_mcp_stream(model, upstream_url, (tool_round(), answer_round())) created: Final = _first_index(events, "response.created") in_progress: Final = _first_index(events, "response.in_progress") @@ -2424,3 +2484,894 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( assert lifecycle_ids[0] == lifecycle_ids[1] assert lifecycle_ids[2] != lifecycle_ids[0] assert not [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] + + +_FIXED_MARGIN: Final = 0.001 + + +def _streamed(body: Mapping[str, object], **terminal: object) -> httpx.Response: + completed: Final = json.loads(json.dumps({**body, **terminal}, default=dict)) + opened: Final = {**completed, "status": "in_progress", "output": [], "usage": None} + items: Final = tuple(enumerate(completed["output"])) + events: Final = ( + {"type": "response.created", "response": opened}, + {"type": "response.in_progress", "response": opened}, + *({"type": "response.output_item.added", "output_index": index, "item": item} for index, item in items), + *({"type": "response.output_item.done", "output_index": index, "item": item} for index, item in items), + {"type": f"response.{completed['status']}", "response": completed}, + ) + return httpx.Response( + 200, + content=_sse(tuple({**event, "sequence_number": number} for number, event in enumerate(events))), + headers=SSE_HEADERS, + ) + + +class _CutAfterTheFirstEvent(httpx.AsyncByteStream): + async def __aiter__(self): + opened: Final = {**json.loads(json.dumps(_ROUND_TWO_BODY, default=dict)), "status": "in_progress"} + yield _sse(({"type": "response.created", "response": opened, "sequence_number": 0},)) + raise httpx.ReadError("connection reset mid-stream") + + +async def _open_stream_auto_execute(logging_obj: Logging) -> object: + return await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=True, + litellm_logging_obj=logging_obj, + ) + + +async def _drain(stream: object) -> tuple[ResponsesAPIStreamingResponse, ...]: + return tuple([event async for event in cast(Any, stream)]) + + +async def _stream_auto_execute(logging_obj: Logging) -> tuple[ResponsesAPIStreamingResponse, ...]: + with _echo_tool_served("PING"): + return await _drain(await _open_stream_auto_execute(logging_obj)) + + +def _final_response( + events: Sequence[ResponsesAPIStreamingResponse], terminal: str = "response.completed" +) -> ResponsesAPIResponse: + assert _event_type(events[-1]) == terminal + response: Final = getattr(events[-1], "response", None) + assert isinstance(response, ResponsesAPIResponse) + return response + + +def _assert_one_stream_row_bills_both_rounds_once( + logged: tuple[StandardLoggingPayload, ...], client: ResponsesAPIResponse +) -> None: + charged_once: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN + assert len(logged) == 1 + row: Final = logged[0] + assert _output_types(row["response"]["output"]) == _output_types(client.output) + assert "message" in _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["cache_hit"] is None + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["hidden_params"]["additional_headers"]["llm_provider-content-type"] == SSE_HEADERS["content-type"] + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (71 + 149, 19 + 5) + assert client.usage.cost == pytest.approx(charged_once) + + +def _assert_breakdown_agrees_with_the_billed_cost(row: StandardLoggingPayload) -> None: + breakdown: Final = row.get("cost_breakdown") + assert breakdown is None or breakdown["total_cost"] == pytest.approx(row["response_cost"]) + + +def _request_failure_row(logging_obj: Logging, success_row: StandardLoggingPayload) -> Mapping[str, object]: + failure_row: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) + failure_payload: Final = cast(StandardLoggingPayload, failure_row["standard_logging_object"]) + assert failure_payload["status"] == "failure" + assert failure_payload["id"] != success_row["id"] + return failure_row + + +def _assert_billed_tokens(row: Mapping[str, object], tokens: tuple[int, int]) -> None: + usage: Final = row["combined_usage_object"] + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens) == tokens + + +def _assert_one_row_bills(logged: tuple[StandardLoggingPayload, ...], body: Mapping[str, object]) -> None: + usage: Final = cast(Mapping[str, int], body["usage"]) + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == ( + usage["input_tokens"], + usage["output_tokens"], + ) + assert logged[0]["response_cost"] == pytest.approx(_round_cost(body)) + assert logged[0]["response_cost"] > 0 + + +def _assert_one_row_bills_round_one_twice_with_one_fixed_margin(logged: tuple[StandardLoggingPayload, ...]) -> None: + charged_once: Final = 2 * _round_cost(_ROUND_ONE_BODY) - _FIXED_MARGIN + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (2 * 71, 2 * 19) + assert logged[0]["response_cost"] == pytest.approx(charged_once) + assert logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_logs_the_final_response_billed_for_both_rounds_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-rounds")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_when_the_proxy_releases( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-rounds-deferred") + + with _echo_tool_served("PING"): + stream: Final = await _open_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + await asyncio.sleep(0.2) + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + assert logged_before_release == () + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_paid_first_round_on_its_success_row_when_the_follow_up_stream_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), httpx.Response(200, stream=_CutAfterTheFirstEvent(), headers=SSE_HEADERS)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-breaks") + + with pytest.raises(litellm.ServiceUnavailableError): + await _stream_auto_execute(logging_obj) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + assert _request_failure_row(logging_obj, recorder.logged[0]).get("response_cost", 0.0) == 0.0 + + +@pytest.mark.asyncio +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-rounds-blocked") + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served("PING"): + stream: Final = await _open_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + if not ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked): + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_paid_first_round_once_when_the_follow_up_is_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + context_window_error: Final = { + "error": { + "message": "Your input exceeds the context window of this model.", + "type": "invalid_request_error", + "param": "input", + "code": "context_length_exceeded", + } + } + _serve_responses( + respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), httpx.Response(400, json=context_window_error)) + ) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-follow-up-rejected")) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert _event_type(events[-1]) == "error" + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_every_paid_round_with_one_fixed_margin_when_the_third_round_is_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + ( + _streamed(_ROUND_ONE_BODY), + _streamed(_ROUND_ONE_BODY), + httpx.Response(400, json={"error": {"message": "Your input exceeds the context window of this model."}}), + ), + ) + recorder: Final = _RoundRecorder() + + await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-third-round-rejected")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one_twice_with_one_fixed_margin(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_every_finished_round_when_the_client_disconnects_during_a_later_tool_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_ONE_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-disconnect-in-tool-call") + echo_then_disconnect: Final = AsyncMock( + side_effect=[[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}], asyncio.CancelledError()] + ) + + with ( + _echo_tool_served("PING"), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", echo_then_disconnect), + pytest.raises(asyncio.CancelledError), + ): + await _drain(await _open_stream_auto_execute(logging_obj)) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one_twice_with_one_fixed_margin(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_TWO_BODY),)) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-one-round")) + await _settle(lambda: len(recorder.logged) > 0) + + assert _output_types(_final_response(events).output) == ("message",) + _assert_one_row_bills(recorder.logged, _ROUND_TWO_BODY) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_both_rounds_with_one_fixed_margin_when_the_final_round_is_incomplete( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + truncated: Final = _streamed( + _ROUND_TWO_BODY, status="incomplete", incomplete_details={"reason": "max_output_tokens"} + ) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), truncated)) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-final-round-incomplete")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events, "response.incomplete")) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_a_failed_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + server_error: Final = {"code": "server_error", "message": "The server had an error processing your request."} + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY, status="failed", error=server_error)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-failed") + + with pytest.raises(litellm.ServiceUnavailableError): + await _stream_auto_execute(logging_obj) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (149, 5) + assert recorder.failed[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(recorder.failed[0]) + failure_row: Final = _request_failure_row(logging_obj, recorder.logged[0]) + _assert_billed_tokens(failure_row, (149, 5)) + assert failure_row["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN) + + +_XAI_MODEL: Final = "xai/grok-4.7" +_XAI_RESPONSES_URL: Final = r"https://api\.x\.ai/v1/responses" + + +def _xai_round(body: Mapping[str, object], cost_in_usd_ticks: int) -> Mapping[str, object]: + usage: Final = cast(Mapping[str, object], body["usage"]) + stated: Final = {**usage, "cost_in_usd_ticks": cost_in_usd_ticks, "context_details": {"cached_tokens": 3}} + return MappingProxyType({**body, "model": "grok-4.7", "usage": stated}) + + +def _xai_round_cost(body: Mapping[str, object]) -> float: + usage: Final = cast(Mapping[str, object], body["usage"]) + stated: Final = {**usage, "cost": xai_reported_cost_in_usd(usage["cost_in_usd_ticks"])} + return litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate({**body, "usage": stated}), + model=_XAI_MODEL, + custom_llm_provider="xai", + call_type="aresponses", + ) + + +@pytest.mark.asyncio +async def test_xai_stream_auto_execute_bills_each_round_at_its_own_price_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": {"fixed_amount": _FIXED_MARGIN}}) + rounds: Final = (_xai_round(_ROUND_ONE_BODY, 12_400_000), _xai_round(_ROUND_TWO_BODY, 8_800_000)) + _serve_responses(respx_mock, monkeypatch, tuple(_streamed(body) for body in rounds), _XAI_RESPONSES_URL) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await litellm.aresponses( + model=_XAI_MODEL, + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="xai-mcp-rounds", + stream=True, + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-xai-stream-rounds"), + ) + ) + await _settle(lambda: len(recorder.logged) > 0) + + charged_once: Final = _xai_round_cost(rounds[0]) + _xai_round_cost(rounds[1]) - _FIXED_MARGIN + client: Final = _final_response(events) + assert len(recorder.logged) == 1 + assert recorder.logged[0]["response_cost"] == pytest.approx(charged_once) + assert recorder.logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert client.usage.cost == pytest.approx(charged_once) + assert json.loads(events[-1].model_dump_json())["response"]["usage"]["context_details"] == {"cached_tokens": 6} + + +_UNPRICED_MODEL: Final = "litellm-unpriced-mcp-rounds" + + +@pytest.mark.asyncio +async def test_stream_auto_execute_on_an_unpriced_model_logs_both_rounds_as_a_failed_cost_lookup_not_a_free_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + rounds: Final = tuple( + MappingProxyType({**body, "model": _UNPRICED_MODEL}) for body in (_ROUND_ONE_BODY, _ROUND_TWO_BODY) + ) + _serve_rounds(respx_mock, monkeypatch, rounds) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await litellm.aresponses( + model=f"openai/{_UNPRICED_MODEL}", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=True, + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-stream-unpriced-rounds"), + ) + ) + await _settle(lambda: len(recorder.logged) > 0) + + client: Final = _final_response(events) + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["response_cost_failure_debug_info"] is not None + assert client.usage is not None + assert client.usage.cost is None + + +async def _stop_reading_once_the_follow_up_answers(logging_obj: Logging) -> None: + stream: Final = await _open_stream_auto_execute(logging_obj) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_item.added" and _output_types((event.item,)) == ("message",): + return + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_finished_first_round_when_the_caller_stops_reading_the_follow_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _stop_reading_once_the_follow_up_answers(_rounds_logging_obj(recorder, "mcp-stream-caller-stops")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + + +async def _stream_read_until_the_follow_up_answers(logging_obj: Logging) -> object: + stream: Final = await _open_stream_auto_execute(logging_obj) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_item.added" and _output_types((event.item,)) == ("message",): + return stream + pytest.fail("the follow-up round never answered") + + +_STREAM_BILLING_NAMES: Final = ("StreamedRoundBilling", "dispatch_success_handlers", "StreamingIterator") + + +def _names_stream_billing(text: str) -> bool: + return any(name in text for name in _STREAM_BILLING_NAMES) + + +class _UnraisableRecorder: + def __init__(self) -> None: + self.raised: tuple[str, ...] = () + + def __call__(self, unraisable: "sys.UnraisableHookArgs") -> None: + self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") + + +def test_stream_auto_execute_dropped_after_its_event_loop_closed_discards_its_parked_log_cleanly( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + unraisable: Final = _UnraisableRecorder() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + + with _echo_tool_served("PING"), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + asyncio.run(_stream_read_until_the_follow_up_answers(_rounds_logging_obj(recorder, "mcp-stream-loop-closed"))) + gc.collect() + + assert recorder.logged == () + assert [raised for raised in unraisable.raised if _names_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_stream_billing(str(warning.message))] == [] + + +_BRIDGED_MODEL: Final = "anthropic/claude-haiku-5-5" +_ANTHROPIC_MESSAGES_URL: Final = r"https://api\.anthropic\.com/v1/messages" + + +def _bridged_round(content: bytes) -> httpx.Response: + return httpx.Response(200, content=content, headers=SSE_HEADERS) + + +def _bridged_echo_call_round(cache_usage: Mapping[str, object] = MappingProxyType({})) -> httpx.Response: + return _bridged_round( + _anthropic_round( + "msg_echo_call", + {"type": "tool_use", "id": "call_echo", "name": "echo", "input": {}}, + {"type": "input_json_delta", "partial_json": json.dumps({"value": "PING"})}, + "tool_use", + model="claude-haiku-5-5", + usage=(533, 40), + cache_usage=cache_usage, + ) + ) + + +def _bridged_echo_answer_round( + stop: str = "end_turn", cache_usage: Mapping[str, object] = MappingProxyType({}) +) -> httpx.Response: + return _bridged_round( + _anthropic_round( + "msg_echo_answer", + {"type": "text", "text": ""}, + {"type": "text_delta", "text": "PING"}, + stop, + model="claude-haiku-5-5", + usage=(470, 4), + cache_usage=cache_usage, + ) + ) + + +def _anthropic_prompt_cache_usage(written: int, read: int) -> Mapping[str, object]: + return MappingProxyType( + { + "cache_creation_input_tokens": written, + "cache_read_input_tokens": read, + "cache_creation": {"ephemeral_5m_input_tokens": written, "ephemeral_1h_input_tokens": 0}, + "inference_geo": "us", + "speed": "standard", + } + ) + + +class _CutAfterTheMessageStart(httpx.AsyncByteStream): + async def __aiter__(self): + yield _sse((_anthropic_message_start("msg_echo_answer", "claude-haiku-5-5", 470),)) + raise httpx.ReadError("connection reset mid-stream") + + +def _bridged_round_cost(prompt_tokens: int, completion_tokens: int) -> float: + return litellm.completion_cost( + completion_response=ModelResponse( + model="claude-haiku-5-5", + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ), + model=_BRIDGED_MODEL, + ) + + +async def _open_bridged_stream_auto_execute(logging_obj: Logging) -> object: + return await litellm.aresponses( + model=_BRIDGED_MODEL, + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-ant-mcp-rounds", + stream=True, + litellm_logging_obj=logging_obj, + ) + + +def _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder: _RoundRecorder, answer: str) -> None: + assert len(recorder.results) == 1 + logged: Final = recorder.results[0] + assert isinstance(logged, ModelResponse) + assert logged.choices[0].message.content == answer + usage: Final[object] = getattr(logged, "usage", None) + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens) == (533 + 470, 40 + 4) + + +def _assert_one_bridged_row_bills_both_rounds_once(recorder: _RoundRecorder, client: ResponsesAPIResponse) -> None: + charged_once: Final = _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) - _FIXED_MARGIN + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder, "PING") + assert row["id"] == client.id + assert "message" in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (533 + 470, 40 + 4) + assert row["cache_hit"] is False + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (533 + 470, 40 + 4) + assert client.usage.cost == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_logs_the_final_response_billed_for_both_rounds_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-stream-rounds")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert _output_types(_final_response(events).output) == ("mcp_call", "message") + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_logs_the_provider_usage_of_both_rounds_summed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + ( + _bridged_echo_call_round(_anthropic_prompt_cache_usage(written=10, read=5)), + _bridged_echo_answer_round(cache_usage=_anthropic_prompt_cache_usage(written=3, read=7)), + ), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-provider-usage")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + usage: Final = row["metadata"]["usage_object"] + assert "choices" in row["response"] + assert (row["prompt_tokens"], row["completion_tokens"]) == (533 + 10 + 5 + 470 + 3 + 7, 40 + 4) + assert (usage["cache_creation_input_tokens"], usage["cache_read_input_tokens"]) == (10 + 3, 5 + 7) + assert usage["prompt_tokens_details"]["cache_creation_token_details"]["ephemeral_5m_input_tokens"] == 10 + 3 + assert (usage["inference_geo"], usage["speed"]) == ("us", "standard") + assert usage["cost"] == pytest.approx(row["response_cost"]) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_both_rounds_with_one_fixed_margin_when_the_final_round_is_incomplete( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), _bridged_echo_answer_round("max_tokens")), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-stream-incomplete")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events, "response.incomplete")) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_on_release( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-rounds-deferred") + + with _echo_tool_served("PING"): + stream: Final = await _open_bridged_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + await asyncio.sleep(0.2) + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + assert logged_before_release == () + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_rounds_on_the_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-rounds-blocked") + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served("PING"): + stream: Final = await _open_bridged_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _drain(stream) + ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await _settle(lambda: len(recorder.failed) > 0) + + charged_once: Final = _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) - _FIXED_MARGIN + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert recorder.failed[0]["response_cost"] == pytest.approx(charged_once) + assert recorder.failed[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stop", ("end_turn", "max_tokens"), ids=("completed", "incomplete")) +async def test_bridged_stream_auto_execute_bills_both_rounds_when_the_caller_stops_reading_after_the_answer_text( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stop: str +): + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round(stop)), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + + async def stop_reading_once_the_answer_text_is_done() -> None: + stream: Final = await _open_bridged_stream_auto_execute( + _rounds_logging_obj(recorder, "mcp-bridged-stream-caller-stops") + ) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_text.done": + return + + with _echo_tool_served("PING"): + await stop_reading_once_the_answer_text_is_done() + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder, "PING") + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert recorder.logged[0]["response_cost"] == pytest.approx( + _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) + ) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), httpx.Response(200, stream=_CutAfterTheMessageStart(), headers=SSE_HEADERS)), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-breaks") + + with _echo_tool_served("PING"), pytest.raises(MidStreamFallbackError): + await _drain(await _open_bridged_stream_auto_execute(logging_obj)) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert len(recorder.logged) == 1 + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (533, 40) + assert recorder.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(533, 40)) + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (470, 0) + assert recorder.failed[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 0) - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(recorder.failed[0]) + failure_row: Final = _request_failure_row(logging_obj, recorder.logged[0]) + _assert_billed_tokens(failure_row, (470, 0)) + assert failure_row["response_cost"] == pytest.approx(_bridged_round_cost(470, 0) - _FIXED_MARGIN) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_a_broken_follow_up_its_own_fixed_margin_after_a_cached_first_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + _serve_responses( + respx_mock, + monkeypatch, + ( + _bridged_echo_call_round(), + _bridged_echo_answer_round(), + httpx.Response(200, stream=_CutAfterTheMessageStart(), headers=SSE_HEADERS), + ), + _ANTHROPIC_MESSAGES_URL, + ) + warm: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(warm, "mcp-bridged-cache-then-break-warm")) + ) + await _settle(lambda: len(warm.logged) > 0) + + recorder: Final = _RoundRecorder() + with _echo_tool_served("PONG"), pytest.raises(MidStreamFallbackError): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-cache-then-break")) + ) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert [row["response_cost"] for row in recorder.logged] == [0.0] + assert len(recorder.failed) == 1 + assert recorder.failed[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 0)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_bridged_echo_answer_round(),), _ANTHROPIC_MESSAGES_URL) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _drain(await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-one-round"))) + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + assert "choices" in recorder.logged[0]["response"] + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (470, 4) + assert recorder.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 4)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_nothing_for_a_round_served_from_the_response_cache( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + provider: Final = _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), _bridged_echo_answer_round(), _bridged_echo_answer_round()), + _ANTHROPIC_MESSAGES_URL, + ) + warm: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain(await _open_bridged_stream_auto_execute(_rounds_logging_obj(warm, "mcp-bridged-cache-warm"))) + await _settle(lambda: len(warm.logged) > 0) + + first_cached: Final = _RoundRecorder() + with _echo_tool_served("PONG"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(first_cached, "mcp-bridged-cache-first-round")) + ) + await _settle(lambda: len(first_cached.logged) > 0) + + both_cached: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(both_cached, "mcp-bridged-cache-both-rounds")) + ) + await _settle(lambda: len(both_cached.logged) > 0) + + client: Final = _final_response(events) + assert provider.call_count == 3 + assert len(first_cached.logged) == 1 + assert (first_cached.logged[0]["prompt_tokens"], first_cached.logged[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert first_cached.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 4)) + assert first_cached.logged[0]["cache_hit"] is False + assert [(row["response_cost"], row["cache_hit"]) for row in both_cached.logged] == [(0.0, True)] + assert client.usage is not None + assert client.usage.cost == pytest.approx(_bridged_round_cost(470, 4)) diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index 8eb73b7c6e1..ed39da1b7a9 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py @@ -2,6 +2,7 @@ import sys import types from contextlib import nullcontext from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -9,7 +10,6 @@ from mcp.types import CallToolResult, TextContent import litellm # noqa: F401 - ensures litellm.responses.main is registered in sys.modules from litellm.responses.mcp.mcp_streaming_iterator import ( - MAX_MCP_TOOL_CALL_ROUNDS, MCPEnhancedStreamingIterator, ) from litellm.types.llms.openai import ( @@ -96,7 +96,7 @@ def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: return call_tool -def _make_iterator(initial_chunks) -> MCPEnhancedStreamingIterator: +def _make_iterator(initial_chunks, max_agentic_loops: int | str | None = None) -> MCPEnhancedStreamingIterator: return MCPEnhancedStreamingIterator( base_iterator=_FakeAsyncStream(initial_chunks), mcp_events=[], @@ -107,6 +107,7 @@ def _make_iterator(initial_chunks) -> MCPEnhancedStreamingIterator: "model": "gpt-4", "input": "what is berriai/litellm?", "tools": [{"type": "mcp"}], + **({"max_agentic_loops": max_agentic_loops} if max_agentic_loops is not None else {}), }, ) @@ -157,35 +158,43 @@ async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeyp @pytest.mark.asyncio -async def test_tool_call_rounds_are_capped(monkeypatch): +@pytest.mark.parametrize( + ("max_agentic_loops", "expected_cap"), + ( + pytest.param(None, 5, id="unset"), + pytest.param(7, 7, id="integer"), + pytest.param("7", 7, id="integer-string"), + pytest.param(2, 2, id="custom-cap"), + ), +) +async def test_tool_call_rounds_are_capped(monkeypatch, max_agentic_loops, expected_cap): """ A model that keeps calling tools every round must not loop forever — - auto-execution stops at MAX_MCP_TOOL_CALL_ROUNDS, and the follow-up made + auto-execution stops at the configured round cap, and the follow-up made at the cap drops "tools" from the request so the model is forced to answer in text instead of the stream just hanging. """ call_tool = _mock_mcp_environment(monkeypatch) - # Rounds 2..MAX_MCP_TOOL_CALL_ROUNDS keep calling the tool; the call made - # once the cap is hit returns a text-only response. - tool_call_streams = [ - _tool_call_stream(f"call_{i}", "read_wiki_contents") for i in range(2, MAX_MCP_TOOL_CALL_ROUNDS + 1) - ] - aresponses_mock = AsyncMock(side_effect=[*tool_call_streams, _text_only_stream("giving up on tools")]) + tool_call_streams: Final = tuple( + _tool_call_stream(f"call_{i}", "read_wiki_contents") for i in range(2, expected_cap + 1) + ) + aresponses_mock: Final = AsyncMock(side_effect=(*tool_call_streams, _text_only_stream("giving up on tools"))) monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) iterator = _make_iterator( [ _output_item_added_chunk(), _completed_chunk([_function_call("call_1", "read_wiki_contents")]), - ] + ], + max_agentic_loops=max_agentic_loops, ) _ = [chunk async for chunk in iterator] - assert iterator.tool_call_round == MAX_MCP_TOOL_CALL_ROUNDS - assert call_tool.call_count == MAX_MCP_TOOL_CALL_ROUNDS - assert aresponses_mock.call_count == MAX_MCP_TOOL_CALL_ROUNDS + assert iterator.tool_call_round == expected_cap + assert call_tool.call_count == expected_cap + assert aresponses_mock.call_count == expected_cap # Every follow-up before the cap still offered tools; only the capped one drops them. for call in aresponses_mock.call_args_list[:-1]: @@ -193,6 +202,12 @@ async def test_tool_call_rounds_are_capped(monkeypatch): assert "tools" not in aresponses_mock.call_args_list[-1].kwargs +@pytest.mark.parametrize(("max_agentic_loops", "expected_error"), ((0, ValueError), ("three", TypeError))) +def test_invalid_max_agentic_loops_raises_when_creating_iterator(max_agentic_loops, expected_error): + with pytest.raises(expected_error, match=r"litellm_params\.max_agentic_loops"): + _make_iterator([], max_agentic_loops=max_agentic_loops) + + @pytest.mark.asyncio async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch): """ diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index 63204114501..278bd4c1e52 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -10,6 +10,7 @@ import copy import json from collections.abc import Mapping from importlib import import_module +from pathlib import Path from typing import Final, cast from unittest.mock import MagicMock, patch @@ -852,3 +853,92 @@ def test_bridge_reasoning_effort_reads_only_a_string_or_mapping_effort_nested_in extra_body: JsonValue, expected: JsonValue ) -> None: assert bridge_reasoning_effort(None, extra_body) == expected + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("prefix", ["", "responses/"]) +@pytest.mark.parametrize("catalog_mode", [None, "chat", "responses"]) +@pytest.mark.asyncio +async def test_chatgpt_completion_routes_to_responses_without_a_provider_catalog_entry( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + respx_mock: respx.MockRouter, + asynchronous: bool, + stream: bool, + prefix: str, + catalog_mode: str | None, +) -> None: + model: Final = "test-chatgpt-routing" + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("CHATGPT_API_BASE", "https://chatgpt.example") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "fake-access-token", "expires_at": 253402300799, "account_id": "fake-account"}) + ) + monkeypatch.delitem(litellm.model_cost, f"chatgpt/{model}", raising=False) + monkeypatch.delitem(litellm.model_cost, model, raising=False) + if catalog_mode is not None: + monkeypatch.setitem(litellm.model_cost, model, {"mode": catalog_mode, "litellm_provider": "openai"}) + litellm.get_model_info.cache_clear() + payload: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1, + "status": "completed", + "model": model, + "output": [ + { + "type": "message", + "id": "msg_test", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + events: Final = ( + {"type": "response.created", "response": {**payload, "status": "in_progress", "output": []}}, + { + "type": "response.output_text.delta", + "delta": "hello", + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + }, + {"type": "response.completed", "response": payload}, + ) + upstream: Final = respx_mock.post("https://chatgpt.example/responses").respond( + 200, + headers={"content-type": "text/event-stream"}, + text="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events), + ) + kwargs: Final = { + "model": f"chatgpt/{prefix}{model}", + "messages": [{"role": "user", "content": "Say hello"}], + "api_base": "https://chatgpt.example", + "stream": stream, + "num_retries": 0, + } + response: Final = await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs) + result: Final = ( + litellm.stream_chunk_builder([chunk async for chunk in response] if asynchronous else list(response)) + if stream + else response + ) + assert result is not None + assert result.choices[0].message.content == "hello" + assert upstream.call_count == 1 + sent: Final = json.loads(upstream.calls[0].request.content) + assert sent["model"] == model + assert sent["store"] is False + assert sent["input"] == [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Say hello"}]} + ] + + +@pytest.mark.parametrize("provider", ["openai", "azure", "xai", "github_copilot"]) +def test_missing_catalog_entry_does_not_force_other_providers_to_responses(provider: str) -> None: + info, model = responses_api_bridge_check(model="test-chatgpt-routing-control", custom_llm_provider=provider) + assert (info.get("mode"), model) == (None, "test-chatgpt-routing-control") diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index ea3eab7e002..bf571c74e22 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -4,9 +4,12 @@ over the wire and surface provider errors correctly. Expected JSON bodies are st in expected_responses_api_request/. """ +import asyncio import copy import json import logging +import math +from collections.abc import Awaitable, Callable, Mapping from pathlib import Path from importlib import import_module from typing import Final @@ -15,10 +18,12 @@ from unittest.mock import AsyncMock, patch import httpx import pytest import respx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.llms.openai import ResponsesAPIResponse from tests.unit.llms.bedrock.slow_upstream import ( STREAM_TIMEOUT_SECONDS, UPSTREAM_FIRST_BYTE_SECONDS, @@ -331,7 +336,9 @@ async def test_aresponses_sends_the_declared_reasoning_effort_level_rounded_up( def _effort_levels_disabled_by(entry: dict) -> tuple[str, ...]: - return tuple(level for level in ("none", "minimal", "low") if entry.get(f"supports_{level}_reasoning_effort") is False) + return tuple( + level for level in ("none", "minimal", "low") if entry.get(f"supports_{level}_reasoning_effort") is False + ) def _gpt_5_row_disabling_a_level(provider: str) -> tuple[str, str]: @@ -621,7 +628,8 @@ async def test_aresponses_websocket_strips_responses_routing_prefix_from_openai_ from litellm.responses.main import _aresponses_websocket with patch.object( - import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket", + import_module("litellm.responses.main").base_llm_http_handler, + "async_responses_websocket", new_callable=AsyncMock, ) as mock_ws: await _aresponses_websocket( @@ -643,7 +651,8 @@ async def test_aresponses_websocket_keeps_routing_hints_out_of_the_relay_kwargs( from litellm.responses.main import _aresponses_websocket with patch.object( - import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket", + import_module("litellm.responses.main").base_llm_http_handler, + "async_responses_websocket", new_callable=AsyncMock, ) as mock_ws: await _aresponses_websocket( @@ -669,7 +678,9 @@ _ORIGINAL_WS_INPUT = [ @pytest.mark.asyncio @pytest.mark.parametrize("nested", [False, True]) -async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame(nested: bool): # test-quality-ok: the first frame handed to the relay is the only place the routed input is observable before the provider socket +async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame( + nested: bool, +): # test-quality-ok: the first frame handed to the relay is the only place the routed input is observable before the provider socket from unittest.mock import MagicMock from litellm.responses.main import _aresponses_websocket @@ -680,7 +691,8 @@ async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame ) with patch.object( - import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket", + import_module("litellm.responses.main").base_llm_http_handler, + "async_responses_websocket", new_callable=AsyncMock, ) as mock_ws: await _aresponses_websocket( @@ -709,7 +721,8 @@ async def test_aresponses_websocket_forwards_the_first_frame_verbatim_when_routi first_message = '{"type": "response.create", "model": "gpt-5.6", "input": [{"role": "user", "content": "hi"}]}' with patch.object( - import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket", + import_module("litellm.responses.main").base_llm_http_handler, + "async_responses_websocket", new_callable=AsyncMock, ) as mock_ws: await _aresponses_websocket( @@ -966,9 +979,7 @@ async def test_injection_points_reach_a_foundry_deployment_of_an_openai_model(mo Azure OpenAI one the Foundry deployment got Anthropic marks the Responses transform then stripped.""" monkeypatch.setenv("AZURE_AI_API_BASE", "https://other-deployment.openai.azure.com") injected_client = AsyncHTTPHandler() - mock_post = AsyncMock( - return_value=MockResponse(_minimal_responses_api_payload("resp_foundry", "gpt-6-astra"), 200) - ) + mock_post = AsyncMock(return_value=MockResponse(_minimal_responses_api_payload("resp_foundry", "gpt-6-astra"), 200)) injected_client.post = mock_post await litellm.aresponses( @@ -1141,6 +1152,61 @@ async def test_a_native_responses_provider_places_every_point_itself(): assert with_unmatchable_system == user_only +_CLAUDE_COMPACTION_ITEM: Final = { + "type": "compaction", + "id": "cmp_claude", + "encrypted_content": json.dumps({"type": "compaction", "content": "The user's cat is named Biscuit."}), +} +_OPENAI_COMPACTION_ITEM: Final = {"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAAABopaque"} +_HISTORY: Final = ( + {"role": "system", "content": "Answer briefly."}, + {"role": "user", "content": "My cat is named Biscuit."}, + {"role": "assistant", "content": "Noted."}, +) +_NEXT_TURN: Final = {"role": "user", "content": "What is my cat's name?"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call", "upstream_url"), + [ + pytest.param(litellm.aresponses, "https://api.openai.com/v1/responses", id="create"), + pytest.param(litellm.acompact_responses, "https://api.openai.com/v1/responses/compact", id="compact"), + ], +) +@pytest.mark.parametrize( + ("compaction_item", "sent_input"), + [ + ( + _CLAUDE_COMPACTION_ITEM, + [ + _HISTORY[0], + {"role": "user", "content": "The user's cat is named Biscuit.", "type": "message"}, + _NEXT_TURN, + ], + ), + (_OPENAI_COMPACTION_ITEM, [*_HISTORY, _OPENAI_COMPACTION_ITEM, _NEXT_TURN]), + ], +) +async def test_aresponses_sends_a_claude_compaction_item_to_openai_as_its_summary( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + call: Callable[..., Awaitable[object]], + upstream_url: str, + compaction_item: dict[str, str], + sent_input: list[dict[str, str]], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(upstream_url).mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_compacted", "gpt-5.6")) + ) + + await call(model="openai/gpt-5.6", api_key="fake-api-key", input=[*_HISTORY, compaction_item, _NEXT_TURN]) + + assert json.loads(upstream.calls[0].request.read())["input"] == sent_input + + def test_non_streamed_response_ignores_a_stream_timeout_it_never_uses() -> None: answer: Final = json.dumps(_minimal_responses_api_payload("resp_unstreamed", "gpt-5.6")).encode() client: Final = HTTPHandler( @@ -1175,3 +1241,227 @@ def test_faked_stream_accepts_a_text_timeout(monkeypatch: pytest.MonkeyPatch) -> assert [event.response.output[0].content[0].text for event in response if event.type == "response.completed"] == [ "Done." ] + + +async def _invoke_responses(sync_mode: bool, model: str, pricing: Mapping[str, float]) -> ResponsesAPIResponse: + response: Final = ( + litellm.responses(model=model, input="hi", **pricing) + if sync_mode + else await litellm.aresponses(model=model, input="hi", **pricing) + ) + assert isinstance(response, ResponsesAPIResponse) + return response + + +def _response_cost(response: ResponsesAPIResponse) -> float: + hidden_params: Final[dict[str, float]] = getattr(response, "_hidden_params") + return hidden_params["response_cost"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("model", ("openai/gpt-5.4-nano", "openai/custom-priced-model-not-in-cost-map")) +async def test_responses_bills_per_call_custom_pricing_from_usage( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, sync_mode: bool, model: str +): + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + model_cost: Final[dict[str, object]] = getattr(litellm, "model_cost") + monkeypatch.setattr(litellm, "model_cost", dict(model_cost)) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_custom_price", model.split("/")[1])) + ) + pricing: Final = {"input_cost_per_token": 1e-5, "output_cost_per_token": 2e-5} + + response: Final = await _invoke_responses(sync_mode, model, pricing) + + assert response.usage is not None + expected: Final = ( + response.usage.input_tokens * pricing["input_cost_per_token"] + + response.usage.output_tokens * pricing["output_cost_per_token"] + ) + assert math.isclose(_response_cost(response), expected) + + +@pytest.mark.asyncio +async def test_concurrent_responses_calls_keep_their_own_per_call_prices( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + model: Final = "openai/gpt-5.4-nano" + model_cost: Final[dict[str, object]] = getattr(litellm, "model_cost") + cost_map_row_before: Final = copy.deepcopy(model_cost.get(model)) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_concurrent", "gpt-5.4-nano")) + ) + rates: Final = ((1e-5, 2e-5), (3e-5, 4e-5)) + + responses: Final = await asyncio.gather( + *( + _invoke_responses(False, model, {"input_cost_per_token": in_rate, "output_cost_per_token": out_rate}) + for in_rate, out_rate in rates + ) + ) + + for response, (in_rate, out_rate) in zip(responses, rates, strict=True): + assert response.usage is not None + assert math.isclose( + _response_cost(response), response.usage.input_tokens * in_rate + response.usage.output_tokens * out_rate + ) + assert model_cost.get(model) == cost_map_row_before + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("api_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/inference/go/openai/v1", None, "conversation-123"), + ("https://opencode.ai/inference/go/openai/v1", "caller-override", "caller-override"), + ("https://other.example/v1", None, None), + ], +) +@pytest.mark.asyncio +async def test_responses_opencode_session_header( + asynchronous: bool, stream: bool, api_base: str, explicit_header: str | None, expected_session: str | None +) -> None: + received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + payload: Final = _minimal_responses_api_payload("resp_test", "gpt-6-luna") + + def respond(request: httpx.Request) -> httpx.Response: + received.put_nowait(request) + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join( + f"event: {kind}\ndata: {json.dumps({'type': kind, 'response': payload})}\n\n" + for kind in ("response.created", "response.completed") + ), + ) + if stream + else httpx.Response(200, json=payload) + ) + + args: Final = { + "model": "openai/gpt-6-luna", + "input": "hello", + "api_base": api_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + "stream": stream, + "extra_headers": {"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + } + if asynchronous: + async_client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + response: Final = await litellm.aresponses(**args, client=async_client) + if stream: + events: Final = tuple([event.type async for event in response]) + assert "response.completed" in events + else: + assert response.output[0].content[0].text == "Done." + finally: + await async_client.close() + else: + sync_client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + try: + result: Final = litellm.responses(**args, client=sync_client) + if stream: + assert "response.completed" in tuple(event.type for event in result) + else: + assert result.output[0].content[0].text == "Done." + finally: + sync_client.close() + sent: Final = received.get_nowait() + assert str(sent.url) == f"{api_base}/responses" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "gpt-6-luna", "input": "hello", "stream": stream} + assert received.empty() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("source", ["OPENAI_BASE_URL", "OPENAI_API_BASE", "litellm.api_base"]) +@pytest.mark.parametrize("explicit_base", [None, "https://other.example/v1"]) +@pytest.mark.asyncio +async def test_responses_opencode_session_uses_effective_destination( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + asynchronous: bool, + source: str, + explicit_base: str | None, +) -> None: + default_base: Final = "https://opencode.ai/inference/go/openai/v1" + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", default_base if source == "litellm.api_base" else None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if source != "litellm.api_base": + monkeypatch.setenv(source, default_base) + endpoint: Final = respx_mock.post(f"{explicit_base or default_base}/responses").respond( + json=_minimal_responses_api_payload("resp_test", "gpt-6-luna") + ) + args: Final = { + "model": "openai/gpt-6-luna", + "input": "hello", + "api_base": explicit_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + } + response: Final = await litellm.aresponses(**args) if asynchronous else litellm.responses(**args) + assert response.output[0].content[0].text == "Done." + assert endpoint.call_count == 1 + sent: Final = endpoint.calls[0].request + assert sent.headers.get("x-opencode-session") == (None if explicit_base else "conversation-123") + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "gpt-6-luna", "input": "hello"} + + +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.asyncio +async def test_router_responses_deployment_pricing_keeps_inherited_cache_rates( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + model_cost: Final[dict[str, object]] = getattr(litellm, "model_cost") + monkeypatch.setattr(litellm, "model_cost", dict(model_cost)) + litellm.in_memory_llm_clients_cache.flush_cache() + rates: Final = {"input_cost_per_token": 1e-5, "output_cost_per_token": 2e-5} + router: Final = litellm.Router( + model_list=[ + { + "model_name": "nano", + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake-api-key", **rates}, + "model_info": {"id": "nano-deployment"}, + } + ] + ) + payload: Final[dict[str, object]] = _minimal_responses_api_payload("resp_router_cache", "gpt-5.4-nano") + usage: Final = _JSON_OBJECT.validate_python(payload["usage"]) + cached_tokens: Final = 8 + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response( + 200, + json={**payload, "usage": {**usage, "input_tokens_details": {"cached_tokens": cached_tokens}}}, + ) + ) + + response: Final = await router.aresponses(model="nano", input="hi") + + assert isinstance(response, ResponsesAPIResponse) + assert response.usage is not None + uncached_input: Final = response.usage.input_tokens - cached_tokens + full_price: Final = ( + response.usage.input_tokens * rates["input_cost_per_token"] + + response.usage.output_tokens * rates["output_cost_per_token"] + ) + floor: Final = ( + uncached_input * rates["input_cost_per_token"] + response.usage.output_tokens * rates["output_cost_per_token"] + ) + assert floor < _response_cost(response) < full_price diff --git a/tests/unit/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py index 6ee2c495705..0863abd1f1e 100644 --- a/tests/unit/responses/test_responses_utils.py +++ b/tests/unit/responses/test_responses_utils.py @@ -1,5 +1,6 @@ from importlib import import_module import base64 +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -7,10 +8,11 @@ import pytest import litellm from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse -from litellm.types.utils import Usage +from litellm.types.utils import ModelInfo, Usage from litellm.utils import _invalidate_model_cost_lowercase_map @@ -413,6 +415,43 @@ class TestResponseAPILoggingUtils: assert result.prompt_tokens_details.cache_creation_tokens == 10059 assert result.prompt_tokens_details.cached_tokens == 0 + @pytest.mark.parametrize("as_model", [False, True]) + def test_transform_response_api_usage_prices_12h_cache_writes(self, as_model: bool) -> None: + usage_data: Final = { + "input_tokens": 5_000, + "output_tokens": 0, + "total_tokens": 5_000, + "input_tokens_details": { + "cache_write_tokens": 5_000, + "cache_write_tokens_details": {"ttl_30m_tokens": 3_000, "ttl_12h_tokens": 2_000}, + }, + } + usage_input: Final = ResponseAPIUsage(**usage_data) if as_model else usage_data + usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_input) + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost": 3e-6, + "cache_creation_input_token_cost_above_12hr": 5e-6, + } + + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cache_write_tokens == 5_000 + assert usage.prompt_tokens_details.cache_creation_token_details is not None + assert usage.prompt_tokens_details.cache_creation_token_details.ephemeral_12h_input_tokens == 2_000 + assert usage.prompt_tokens_details.cache_creation_token_details.ephemeral_30m_input_tokens == 3_000 + cost: Final = generic_cost_per_token( + model="test-12h-cache-write", + usage=usage, + custom_llm_provider="openai", + model_info=model_info, + )[0] + + assert cost == pytest.approx( + 3_000 * model_info["cache_creation_input_token_cost"] + + 2_000 * model_info["cache_creation_input_token_cost_above_12hr"] + ) + def test_transform_response_api_usage_mixed_details(self): """Test transformation handles mixed token details (cached + image + audio).""" # Setup - hypothetical usage with mixed token types @@ -894,6 +933,42 @@ class TestResponsesInputToChatMessages: ) == [user_message] +class TestMergeProviderSpecificHeaders: + def test_scoped_headers_merge_only_for_matching_provider(self): + scoped = {"custom_llm_provider": "chatgpt", "extra_headers": {"Authorization": "Bearer client-token"}} + + merged = ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers={"originator": "codex_cli_rs"}, + provider_specific_header=scoped, + custom_llm_provider="chatgpt", + ) + untouched = ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers={"originator": "codex_cli_rs"}, + provider_specific_header=scoped, + custom_llm_provider="openai", + ) + + assert merged == {"originator": "codex_cli_rs", "Authorization": "Bearer client-token"} + assert untouched == {"originator": "codex_cli_rs"} + + def test_sequence_of_scoped_entries_each_keep_their_own_scope(self): + entries = [ + {"custom_llm_provider": "anthropic", "extra_headers": {"Authorization": "Bearer anthropic-token"}}, + {"custom_llm_provider": "chatgpt", "extra_headers": {"Authorization": "Bearer chatgpt-token"}}, + ] + + merged = ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers=None, provider_specific_header=entries, custom_llm_provider="chatgpt" + ) + + assert merged == {"Authorization": "Bearer chatgpt-token"} + + def test_no_scoped_headers_returns_extra_headers_unchanged(self): + assert ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers=None, provider_specific_header=None, custom_llm_provider="chatgpt" + ) is None + + def test_responses_object_bills_an_explicit_cache_hit_at_the_explicit_rate(): model = "dashscope/qwen-responses-explicit-cache-test" previous_row = litellm.model_cost.get(model) diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 62d89a55e85..6c3e6177503 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -2605,6 +2605,27 @@ def test_faked_stream_builds_each_refusal_summary_and_custom_tool_input_once_for } +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_usage_estimate_counts_a_compaction_item_as_its_summary(): + assert _estimate_usage_from_text("gpt-4o-mini", _COMPACTED_INPUT, {}, "Dotonbori.") == _estimate_usage_from_text( + "gpt-4o-mini", _SUMMARY_INPUT, {}, "Dotonbori." + ) + + @pytest.mark.asyncio async def test_shutdown_drain_waits_for_a_pending_logging_task(): release: Final = asyncio.Event() diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 04d2b9ddf81..7445070bd3a 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -23,7 +23,8 @@ import asyncio import json from datetime import datetime from importlib import import_module -from typing import Final +from collections.abc import Mapping +from typing import Final, cast from unittest.mock import Mock, patch import httpx @@ -31,6 +32,7 @@ import pytest import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger 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 @@ -44,6 +46,7 @@ from litellm.responses.streaming_iterator import ( _status_code_for_error_fields, ) from litellm.utils import ProviderConfigManager +from litellm.types.utils import StandardLoggingPayload from litellm.types.llms.openai import ( ErrorEvent, ErrorEventError, @@ -523,6 +526,82 @@ def test_handle_logging_failed_response_records_usage_and_cost(): iterator.logging_obj.response_cost_calculator.assert_called_once_with(result=chunk.response) +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failed: tuple[StandardLoggingPayload, ...] = () + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + + +async def _first_failure_row(recorder: _FailureRecorder) -> StandardLoggingPayload: + for _ in range(200): + if recorder.failed: + return recorder.failed[0] + await asyncio.sleep(0.01) + pytest.fail("the failed stream logged no failure row") + + +@pytest.mark.asyncio +async def test_failed_stream_bills_only_its_own_usage_over_an_earlier_attempts_stash(): + recorder: Final = _FailureRecorder() + logging_obj: Final = LiteLLMLoggingObj( + model="gpt-5.6", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime(2026, 1, 1), + litellm_call_id="responses-failed-over-stash", + function_id="responses-failed-over-stash", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.6", optional_params={}, litellm_params=get_litellm_params(custom_llm_provider="openai") + ) + logging_obj.model_call_details["combined_usage_object"] = litellm.Usage( + prompt_tokens=100, completion_tokens=50, total_tokens=150 + ) + logging_obj.model_call_details["response_cost"] = 0.5 + failed: Final = { + "type": "response.failed", + "sequence_number": 1, + "response": { + "id": "resp_failed_over_stash", + "object": "response", + "created_at": 1790000000, + "status": "failed", + "model": "gpt-5.6", + "output": [], + "error": {"code": "server_error", "message": "boom"}, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + } + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=f"data: {json.dumps(failed)}\n\n".encode() + ), + model="gpt-5.6", + responses_api_provider_config=ProviderConfigManager.get_provider_responses_api_config( + provider="openai", model="gpt-5.6" + ), + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + with pytest.raises(MidStreamFallbackError): + await anext(aiter(iterator)) + + row: Final = await _first_failure_row(recorder) + prompt_cost, completion_cost = litellm.cost_per_token( + model="gpt-5.6", prompt_tokens=10, completion_tokens=5, custom_llm_provider="openai" + ) + assert (row["prompt_tokens"], row["completion_tokens"]) == (10, 5) + assert row["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + def test_handle_logging_failed_response_without_usage_skips_recording(): iterator = _make_iterator() iterator.completed_response = _make_failed_chunk( diff --git a/tests/unit/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py index ad61f43c5a0..0ed22658144 100644 --- a/tests/unit/router_strategy/adaptive_router/test_hooks.py +++ b/tests/unit/router_strategy/adaptive_router/test_hooks.py @@ -10,6 +10,7 @@ from litellm.router_strategy.adaptive_router.config import ( ) from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, _recent_tool_results, _resolve_session_key, ) @@ -337,8 +338,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): """The header hook returns the `x-litellm-adaptive-router-model` header so proxy header construction picks it up (works for both streaming and non-streaming; `async_post_call_success_hook` is too late for streaming).""" - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"adaptive_router_chosen_model": "smart"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -348,8 +348,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"litellm_session_id": "sess-A"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -359,8 +358,7 @@ async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_no_metadata(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -370,8 +368,7 @@ async def test_post_call_response_headers_hook_noop_when_no_metadata(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_not_dict(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": "not-a-dict"}, user_api_key_dict=MagicMock(), response=MagicMock(), diff --git a/tests/unit/router_strategy/test_base_routing_strategy.py b/tests/unit/router_strategy/test_base_routing_strategy.py index dc75c3d2916..19ff11dd3f6 100644 --- a/tests/unit/router_strategy/test_base_routing_strategy.py +++ b/tests/unit/router_strategy/test_base_routing_strategy.py @@ -1,15 +1,20 @@ import json import logging -from typing import Any, Dict, List, Optional, Set, Union +from typing import Any, Dict, Final, List, Optional, Set, Union import pytest import asyncio +import threading from unittest.mock import MagicMock, patch +from fakeredis import TcpFakeServer + from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy @@ -162,3 +167,24 @@ async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_erro assert caplog.records == [] assert base_strategy.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_an_increment_counted_after_retire_still_reaches_redis(): + server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis") + threading.Thread(target=server.serve_forever, daemon=True).start() + redis_cache: Final = RedisCache(host="127.0.0.1", port=server.server_address[1]) + strategy: Final = BaseRoutingStrategy( + dual_cache=DualCache(redis_cache=redis_cache, in_memory_cache=InMemoryCache()), + should_batch_redis_writes=True, + default_sync_interval=60, + ) + try: + strategy.retire() + await strategy._increment_value_in_current_window("counted-after-retire", 3, 60) + await asyncio.sleep(0.2) + + assert float(await redis_cache.async_get_cache("counted-after-retire")) == 3 + finally: + server.shutdown() + server.server_close() diff --git a/tests/unit/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py index 4667ac67efa..10704a29347 100644 --- a/tests/unit/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/unit/router_strategy/test_budget_limiter_hotpath.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import gc import logging from types import SimpleNamespace @@ -492,6 +493,8 @@ def _new_router_budget_limiter( budget_limiter._redis_increment_queue_lock = queue_lock if queue_lock is not None else asyncio.Lock() budget_limiter._redis_increment_flush_lock = asyncio.Lock() budget_limiter._detached_increment_operations = None + budget_limiter._retired = False + budget_limiter._flush_after_retire = None return budget_limiter @@ -784,3 +787,57 @@ async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations: await limiter._push_in_memory_increments_to_redis() assert redis_cache.values[_SPEND_KEY] == 10.0 assert limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_queued_when_a_limiter_is_retired_reaches_redis_after_a_failed_write() -> None: + pipeline_started: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + pipeline_completed=pipeline_completed, + should_fail_pipeline=True, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(asyncio.sleep(3600)) + + budget_limiter.retire() + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert redis_cache.values == {} + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_in_flight_when_a_limiter_is_retired_reaches_redis_after_that_write_fails() -> None: + pipeline_started: Final = asyncio.Event() + allow_pipeline_to_complete: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + pipeline_completed=pipeline_completed, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(budget_limiter.periodic_sync_in_memory_spend_with_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + + budget_limiter.retire() + allow_pipeline_to_complete.set() + with contextlib.suppress(asyncio.CancelledError): + await budget_limiter._sync_task + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == [] diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index f956bb66e6c..a46f415e10a 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -112,6 +112,13 @@ def _heuristic_v2_artifact() -> TrainedTierArtifact: ) +def _answer_deployments_named(router: MagicMock) -> None: + """Make a mock router answer `deployments_named` from its own `model_list`, the way Router answers from its index.""" + router.deployments_named.side_effect = lambda name: ( + tuple(deployment for deployment in router.model_list if deployment["model_name"] == name) or None + ) + + @pytest.fixture def mock_router_instance(): """Create a mock LiteLLM Router instance.""" @@ -1065,6 +1072,11 @@ class TestCustomDimensions: "create table widgets (id integer); ALTER TABLE widgets ADD label text;", id="regex", ), + pytest.param( + {"patterns": [r"\b(?Pcreate|alter)\s{1,4}?table\b"]}, + "CREATE TABLE widgets (id integer)", + id="named-group-alternation-lazy-repeat", + ), ], ) def test_custom_dimension_changes_only_matching_requests( @@ -3176,7 +3188,7 @@ class TestCapabilityClassifier: {"model_name": "efficient-model", "litellm_params": {"input_cost_per_token": 0.000001}}, {"model_name": "capable-model", "litellm_params": {"input_cost_per_token": 0.00001}}, ] - mock_router_instance.model_name_to_deployment_indices = {"efficient-model": [0], "capable-model": [1]} + _answer_deployments_named(mock_router_instance) router: Final = self._router( mock_router_instance, adaptive=True, @@ -5665,7 +5677,7 @@ class TestAdaptiveSoftFloors: "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, }, ] - router.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]} + _answer_deployments_named(router) return router @pytest.fixture @@ -6824,7 +6836,7 @@ class TestTierModelAffinity: ) -> ComplexityRouter: mock_router_instance.cache = DualCache() mock_router_instance.model_list = [] - mock_router_instance.model_name_to_deployment_indices = {} + _answer_deployments_named(mock_router_instance) return ComplexityRouter( model_name="affinity-router", litellm_router_instance=mock_router_instance, @@ -7542,7 +7554,7 @@ class TestSessionAffinity: "model_info": {}, }, ] - mock_router_instance.model_name_to_deployment_indices = {"cheap": [0]} + _answer_deployments_named(mock_router_instance) router = ComplexityRouter( model_name="hybrid", litellm_router_instance=mock_router_instance, @@ -11346,7 +11358,7 @@ class TestClientHousekeepingCalls: "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, }, ] - adaptive_instance.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]} + _answer_deployments_named(adaptive_instance) router = ComplexityRouter( model_name="hybrid", litellm_router_instance=adaptive_instance, @@ -13437,7 +13449,7 @@ class TestPlanModeTierFloor: "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, }, ] - adaptive_instance.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]} + _answer_deployments_named(adaptive_instance) router = ComplexityRouter( model_name="hybrid", litellm_router_instance=adaptive_instance, @@ -14936,7 +14948,7 @@ class TestModalityRouting: elif path == "adaptive_pick_rewritten": config["adaptive"] = True mock_router_instance.model_list = [] - mock_router_instance.model_name_to_deployment_indices = {} + _answer_deployments_named(mock_router_instance) router = self._router(mock_router_instance, config, vision) result = await router.async_pre_routing_hook(model="m", request_kwargs=request_kwargs, messages=messages) assert result.model == expected_model diff --git a/tests/unit/router_strategy/test_lowest_latency.py b/tests/unit/router_strategy/test_lowest_latency.py index 0a3a92691ba..706ec567710 100644 --- a/tests/unit/router_strategy/test_lowest_latency.py +++ b/tests/unit/router_strategy/test_lowest_latency.py @@ -17,8 +17,10 @@ from pydantic import ValidationError import litellm from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.responses.utils import ResponseAPILoggingUtils from litellm.router import Router from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler, RoutingArgs +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse DEPLOYMENT_ID = "9876" KWARGS = { @@ -43,7 +45,7 @@ def _embedding_response() -> litellm.EmbeddingResponse: def _recorded_latencies(cache: DualCache): cached = cache.get_cache(key="gemini-embedding-001_map") or {} - return cached.get(DEPLOYMENT_ID, {}).get("latency", []) + return cached.get(DEPLOYMENT_ID, {}).get("latency_normalized", []) def test_sync_embedding_latency_is_json_serializable(): @@ -69,7 +71,7 @@ def test_sync_embedding_latency_is_json_serializable(): ) assert latencies[-1] == pytest.approx(2.0) # the exact failure mode from production: redis cache sync json.dumps - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) @pytest.mark.asyncio @@ -94,7 +96,7 @@ async def test_async_embedding_latency_is_json_serializable(): f"raw timedelta leaked into latency list: {latencies}" ) assert latencies[-1] == pytest.approx(3.0) - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) def _chat_response(completion_tokens: int) -> litellm.ModelResponse: @@ -131,7 +133,7 @@ async def test_async_chat_latency_normalized_per_token(): latencies = _recorded_latencies(cache) assert latencies and latencies[-1] == pytest.approx(0.5) # 2s / 4 tokens - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) @pytest.mark.asyncio @@ -151,7 +153,7 @@ async def test_async_chat_zero_completion_tokens_falls_back_to_seconds(): latencies = _recorded_latencies(cache) assert latencies and latencies[-1] == pytest.approx(3.0) assert not isinstance(latencies[-1], timedelta) - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): @@ -168,12 +170,14 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): latencies = _recorded_latencies(cache) assert latencies and latencies[-1] == pytest.approx(2.0) assert not isinstance(latencies[-1], timedelta) - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) MODEL_GROUP = "gpt-4o-mini" FAST_TTFT_ID = "fast-ttft-short-output" SLOW_TTFT_ID = "slow-ttft-long-output" +FAST_API_BASE = "https://fast.example.test/v1" +SLOW_API_BASE = "https://slow.example.test/v1" STREAMING_DEPLOYMENTS = [ {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {}}, {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {}}, @@ -362,8 +366,8 @@ async def test_ttft_percentile_does_not_change_non_streaming_routing(): cache.set_cache( key=f"{MODEL_GROUP}_map", value={ - FAST_TTFT_ID: {"latency": [1.0], "time_to_first_token_seconds": [0.1]}, - SLOW_TTFT_ID: {"latency": [0.2], "time_to_first_token_seconds": [1.5]}, + FAST_TTFT_ID: {"latency_normalized": [1.0], "time_to_first_token_seconds": [0.1]}, + SLOW_TTFT_ID: {"latency_normalized": [0.2], "time_to_first_token_seconds": [1.5]}, }, ) @@ -380,7 +384,7 @@ async def test_ttft_percentile_does_not_change_non_streaming_routing(): @pytest.mark.asyncio @pytest.mark.parametrize( "cached_entry", - [{"latency": []}, {"2026-09-05-15-39": {"tpm": 28, "rpm": 1}}], + [{"latency_normalized": []}, {"2026-09-05-15-39": {"tpm": 28, "rpm": 1}}], ids=["empty_latency_list", "minute_bucket_only_as_cost_based_routing_writes"], ) async def test_async_get_available_deployments_treats_missing_samples_as_zero_latency(cached_entry): @@ -388,7 +392,7 @@ async def test_async_get_available_deployments_treats_missing_samples_as_zero_la handler = LowestLatencyLoggingHandler(router_cache=cache) cache.set_cache( key="gemini-embedding-001_map", - value={DEPLOYMENT_ID: cached_entry, "slower": {"latency": [0.5]}}, + value={DEPLOYMENT_ID: cached_entry, "slower": {"latency_normalized": [0.5]}}, ) healthy_deployments = [ {"model_info": {"id": DEPLOYMENT_ID}, "litellm_params": {}}, @@ -405,6 +409,69 @@ async def test_async_get_available_deployments_treats_missing_samples_as_zero_la assert picked["model_info"]["id"] == DEPLOYMENT_ID +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("ttft_samples", "expected_latency"), + [([-602.0], 0.0), ([10.0, 10.0, -40.0], 10.0), ([float("nan"), 2.0], 2.0)], + ids=["negative", "negative_average", "nan"], +) +async def test_unusable_ttft_samples_do_not_leave_the_buffer_window_empty( + ttft_samples: list[float], expected_latency: float +): + """``completion_start_time - start_time`` can be negative or NaN. With a non-zero + ``lowest_latency_buffer`` a negative average pushes ``lowest_latency + buffer`` + below every candidate, so the group must still route instead of raising + IndexError from ``random.choice``, and the bad samples must not reach the ranking.""" + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"lowest_latency_buffer": 5}) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={FAST_TTFT_ID: {"time_to_first_token_seconds": ttft_samples}}, + ) + metadata: Final[dict[str, object]] = {} + request_kwargs: Final[dict[str, object]] = {"stream": True, "metadata": metadata} + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=[{"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {"api_base": FAST_API_BASE}}], + request_kwargs=request_kwargs, + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + assert metadata["_latency_per_deployment"] == {FAST_API_BASE: expected_latency} + + +@pytest.mark.asyncio +async def test_unusable_ttft_samples_fall_back_to_recorded_latency(): + """A deployment whose TTFT samples are all unusable is ranked by its recorded + latency, so dirty samples cannot make it look like the fastest deployment.""" + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"lowest_latency_buffer": 5}) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"latency_normalized": [0.1], "time_to_first_token_seconds": [-600.0]}, + SLOW_TTFT_ID: {"latency_normalized": [5.0], "time_to_first_token_seconds": [-300.0]}, + }, + ) + metadata: Final[dict[str, object]] = {} + request_kwargs: Final[dict[str, object]] = {"stream": True, "metadata": metadata} + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=[ + {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {"api_base": FAST_API_BASE}}, + {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {"api_base": SLOW_API_BASE}}, + ], + request_kwargs=request_kwargs, + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + assert metadata["_latency_per_deployment"] == {FAST_API_BASE: 0.1, SLOW_API_BASE: 5.0} + + def _latency_router(routing_strategy_args: dict) -> Router: return Router( model_list=[ @@ -605,7 +672,7 @@ def test_latency_updated(): end_time=end_time, ) latency_key = f"{model_group}_map" - assert end_time - start_time == test_cache.get_cache(key=latency_key)[deployment_id]["latency"][0] + assert end_time - start_time == test_cache.get_cache(key=latency_key)[deployment_id]["latency_normalized"][0] @pytest.mark.usefixtures("frozen_latency_clock") @@ -1002,7 +1069,7 @@ def test_latency_list_trimming_discards_oldest_entry(): latency_key = f"{model_group}_map" cached_data = test_cache.get_cache(key=latency_key) - latency_list = cached_data[deployment_id]["latency"] + latency_list = cached_data[deployment_id]["latency_normalized"] assert len(latency_list) == max_size, f"Expected {max_size} entries, got {len(latency_list)}" @@ -1062,7 +1129,7 @@ async def test_latency_list_trimming_discards_oldest_entry_async(): latency_key = f"{model_group}_map" cached_data = await test_cache.async_get_cache(key=latency_key) - latency_list = cached_data[deployment_id]["latency"] + latency_list = cached_data[deployment_id]["latency_normalized"] assert len(latency_list) == max_size @@ -1129,7 +1196,7 @@ async def test_timeout_penalty_discards_oldest_entry(): latency_key = f"{model_group}_map" cached_data = await test_cache.async_get_cache(key=latency_key) - latency_list = cached_data[deployment_id]["latency"] + latency_list = cached_data[deployment_id]["latency_normalized"] assert len(latency_list) == max_size @@ -1181,7 +1248,7 @@ def test_list_order_preserved_after_multiple_trims(): latency_key = f"{model_group}_map" cached_data = test_cache.get_cache(key=latency_key) - latency_list = cached_data[deployment_id]["latency"] + latency_list = cached_data[deployment_id]["latency_normalized"] assert len(latency_list) == max_size @@ -1423,7 +1490,7 @@ def test_latency_cache_honors_custom_ttl() -> None: latency_key: Final = "ttl-group_map" cached: Final = cache.get_cache(key=latency_key) assert isinstance(cached, dict) - assert cached["deployment"]["latency"] == [1.0] + assert cached["deployment"]["latency_normalized"] == [1.0] assert cache.in_memory_cache.ttl_dict[latency_key] == 103.0 clock.return_value = 104.0 @@ -1461,3 +1528,191 @@ async def test_router_model_group_usage_increases_by_logged_tokens() -> None: updated_usage: Final = await router.get_model_group_usage(model_group=model_group) assert updated_usage[0] == initial_usage[0] + 15 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_callback", [False, True]) +@pytest.mark.parametrize("normalized_usage", [False, True]) +@pytest.mark.parametrize("native_seconds", [2, 3]) +async def test_responses_and_chat_latency_use_comparable_token_units( + async_callback: bool, normalized_usage: bool, native_seconds: int +) -> None: + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache) + usage: Final = ResponseAPIUsage(input_tokens=10, output_tokens=100, total_tokens=110) + native: Final = ResponsesAPIResponse(id="resp-latency", created_at=0, output=[], usage=usage) + logged_native: Final = ( + native.model_copy(update={"usage": ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage)}) + if normalized_usage + else native + ) + original_usage: Final = logged_native.usage + assert original_usage is not None + original_fields: Final = original_usage.model_dump() + start: Final = datetime(2026, 1, 1) + for deployment_id, response, seconds in ( + ("native", logged_native, native_seconds), + ("bridge", _chat_response(100), 3), + ): + kwargs: Final = { + "call_type": "aresponses", + "litellm_params": {"metadata": {"model_group": "mixed"}, "model_info": {"id": deployment_id}}, + } + if async_callback: + await handler.async_log_success_event(kwargs, response, start, start + timedelta(seconds=seconds)) + else: + handler.log_success_event(kwargs, response, start, start + timedelta(seconds=seconds)) + recorded: Final = cache.get_cache(key="mixed_map") + assert recorded["native"]["latency_normalized"] == pytest.approx([native_seconds / 100]) + assert recorded["bridge"]["latency_normalized"] == pytest.approx([3 / 100]) + assert logged_native.usage is original_usage + assert logged_native.usage.model_dump() == original_fields + assert native.usage is usage + deployments: Final = [ + {"model_name": "mixed", "model_info": {"id": name}, "litellm_params": {}} for name in ("native", "bridge") + ] + picked: Final = handler.get_available_deployments( + model_group="mixed", healthy_deployments=deployments, request_kwargs={"stream": False, "metadata": {}} + ) + assert picked["model_info"]["id"] in ({"native"} if native_seconds < 3 else {"native", "bridge"}) + native_counts: Final = tuple(value for key, value in recorded["native"].items() if key != "latency_normalized") + assert native_counts == ({"tpm": 0, "rpm": 1},) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_callback", [False, True]) +@pytest.mark.parametrize( + "response", + [ + litellm.ModelResponse(), + ResponsesAPIResponse(id="resp-missing", created_at=0, output=[], usage=None), + ResponsesAPIResponse( + id="resp-zero", created_at=0, output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=0, total_tokens=10), + ), + ], +) +async def test_responses_without_output_tokens_keep_elapsed_seconds( + async_callback: bool, response: litellm.ModelResponse | ResponsesAPIResponse +) -> None: + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache) + original_usage: Final = getattr(response, "usage", None) + if async_callback: + await handler.async_log_success_event(KWARGS, response, 1.0, 3.0) + else: + handler.log_success_event(KWARGS, response, 1.0, 3.0) + assert _recorded_latencies(cache) == [2.0] + assert getattr(response, "usage", None) is original_usage + + +@pytest.mark.usefixtures("frozen_latency_clock") +@pytest.mark.asyncio +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("preserve_new_samples", [False, True]) +async def test_legacy_latency_does_not_starve_deployments_after_upgrade( + async_mode: bool, preserve_new_samples: bool +) -> None: + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache) + minute: Final = FROZEN_NOW.strftime("%Y-%m-%d-%H-%M") + legacy: Final = { + "native": {"latency": [2.0], "time_to_first_token_seconds": [0.1], minute: {"tpm": 0, "rpm": 1}}, + "bridge": {"latency": [0.03], "time_to_first_token_seconds": [0.2], minute: {"tpm": 110, "rpm": 1}}, + } + cache.set_cache(key="mixed_map", value=copy.deepcopy(legacy)) + deployments: Final = [ + {"model_name": "mixed", "model_info": {"id": name}, "litellm_params": {}} for name in ("native", "bridge") + ] + native: Final = ResponsesAPIResponse( + id="resp-upgrade", created_at=0, output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=100, total_tokens=110), + ) + for iteration in range(2): + bridge_kwargs: Final = { + "litellm_params": {"metadata": {"model_group": "mixed"}, "model_info": {"id": "bridge"}}, + } + if async_mode: + await handler.async_log_success_event(bridge_kwargs, _chat_response(100), 0.0, 3.0) + else: + handler.log_success_event(bridge_kwargs, _chat_response(100), 0.0, 3.0) + selected: Final = ( + await handler.async_get_available_deployments(model_group="mixed", healthy_deployments=deployments) + if async_mode + else handler.get_available_deployments(model_group="mixed", healthy_deployments=deployments) + ) + assert selected["model_info"]["id"] == "native" + native_kwargs: Final = { + "litellm_params": {"metadata": {"model_group": "mixed"}, "model_info": {"id": "native"}}, + } + if async_mode: + await handler.async_log_success_event(native_kwargs, native, 0.0, 2.0) + else: + handler.log_success_event(native_kwargs, native, 0.0, 2.0) + recorded: Final = copy.deepcopy(cache.get_cache(key="mixed_map")) + assert recorded["native"]["latency_normalized"] == pytest.approx( + [0.02] * (iteration + 1 if preserve_new_samples else 1) + ) + assert recorded["native"]["latency"] == [2.0] + assert recorded["native"]["time_to_first_token_seconds"] == [0.1] + assert recorded["bridge"]["time_to_first_token_seconds"] == [0.2] + assert recorded["native"][minute] == {"tpm": 0, "rpm": iteration + 2} + assert recorded["bridge"][minute] == {"tpm": 110 + (iteration + 1) * 110, "rpm": iteration + 2} + streaming: Final = handler.get_available_deployments( + model_group="mixed", healthy_deployments=deployments, request_kwargs={"stream": True} + ) + assert streaming["model_info"]["id"] == "native" + cache.set_cache( + key="mixed_map", + value={ + name: {**(recorded[name] if preserve_new_samples else legacy[name]), minute: recorded[name][minute]} + for name in ("native", "bridge") + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_mode", [False, True]) +async def test_normalized_latency_window_evicts_samples_without_changing_legacy_data(async_mode: bool) -> None: + cache: Final = DualCache() + cache.set_cache(key="gemini-embedding-001_map", value={DEPLOYMENT_ID: {"latency": [2.0]}}) + handler: Final = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": 2}) + for elapsed, expected in ((2.0, [0.02]), (3.0, [0.02, 0.03]), (4.0, [0.03, 0.04])): + if async_mode: + await handler.async_log_success_event(KWARGS, _chat_response(100), 0.0, elapsed) + else: + handler.log_success_event(KWARGS, _chat_response(100), 0.0, elapsed) + recorded: Final = cache.get_cache(key="gemini-embedding-001_map")[DEPLOYMENT_ID] + assert recorded["latency_normalized"] == pytest.approx(expected) + assert recorded["latency"] == [2.0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("existing_samples", [(), (0.02, 0.03)]) +async def test_timeout_penalty_updates_normalized_latency_and_avoids_failed_deployment( + existing_samples: tuple[float, ...], +) -> None: + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": 2}) + cache.set_cache( + key="mixed_map", + value={ + "native": {"latency": [0.001], "latency_normalized": list(existing_samples)}, + "bridge": {"latency": [2.0], "latency_normalized": [0.04]}, + }, + ) + await handler.async_log_failure_event( + kwargs={ + "exception": litellm.Timeout(message="request timed out", model="test", llm_provider="openai"), + "litellm_params": {"metadata": {"model_group": "mixed"}, "model_info": {"id": "native"}}, + }, + response_obj=None, start_time=0.0, end_time=3.0, + ) + recorded: Final = cache.get_cache(key="mixed_map")["native"] + assert recorded["latency_normalized"] == [*existing_samples[-1:], 1000.0] + assert recorded["latency"] == [0.001] + selected: Final = await handler.async_get_available_deployments( + model_group="mixed", + healthy_deployments=[{"model_info": {"id": name}, "litellm_params": {}} for name in ("native", "bridge")], + ) + assert selected["model_info"]["id"] == "bridge" diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 7e4b470f1c1..7235333f59f 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -1249,7 +1249,7 @@ async def test_latency_based_override_learns_from_the_overriding_requests(): ) def samples(deployment_id: str) -> list[float]: - recorded = (router.cache.get_cache("grp_map") or {}).get(deployment_id, {}).get("latency", []) + recorded = (router.cache.get_cache("grp_map") or {}).get(deployment_id, {}).get("latency_normalized", []) return [latency for latency in recorded if latency > 0] async def overriding_call() -> str: @@ -1316,7 +1316,7 @@ async def test_override_matching_a_routing_group_strategy_records_each_request_o deployment_id = response._hidden_params["model_id"] await _async_wait_until(lambda: (router.cache.get_cache("grp_map") or {}).get(deployment_id) is not None) - assert len(router.cache.get_cache("grp_map")[deployment_id]["latency"]) == 1 + assert len(router.cache.get_cache("grp_map")[deployment_id]["latency_normalized"]) == 1 assert request["litellm_logging_obj"].dynamic_success_callbacks is None diff --git a/tests/unit/router_strategy/test_savings_baseline.py b/tests/unit/router_strategy/test_savings_baseline.py index 5efc73d2dcb..4241a541202 100644 --- a/tests/unit/router_strategy/test_savings_baseline.py +++ b/tests/unit/router_strategy/test_savings_baseline.py @@ -107,8 +107,7 @@ class TestResolveBaseline: served; a dashboard counterfactual must not be able to take routing down.""" class Exploding: - @property - def model_name_to_deployment_indices(self): + def deployments_named(self, model_name): raise RuntimeError("router is mid-reload") assert resolve_baseline(Exploding(), ["anything"]) is None diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 9bbaed0ae1b..6b4ea74894a 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,20 +1,368 @@ import asyncio import json from typing import Final +from collections.abc import AsyncIterator, Mapping +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeConnection +from redis.asyncio import BlockingConnectionPool, Redis import litellm from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) +def _repinning_router(enabled: bool | None = True) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "repin-group", + "litellm_params": {"model": "openai/test-model", "api_key": "fake-api-key"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("original", "replacement", "other") + ] + + [ + { + "model_name": "unrelated-group", + "litellm_params": {"model": "openai/test-model", "api_key": "fake-api-key"}, + "model_info": {"id": "unrelated"}, + } + ], + optional_pre_call_checks=["session_affinity", "deployment_affinity"], + **({"session_affinity_repin_on_failover": enabled} if enabled is not None else {}), + ) + + +def _repin_params(deployment_id: str, caller: str = "caller-1") -> dict[str, object]: + return { + "model_info": {"id": deployment_id}, + "metadata": { + "model_group": "repin-group", + "deployment_model_name": "repin-group", + "session_id": "repin-session", + "user_api_key_hash": caller, + }, + } + + +def _repin_key(caller: str = "caller-1") -> str: + return DeploymentAffinityCheck.get_session_affinity_cache_key("repin-group", "repin-session", caller) + + +def _cool_down_original(router: litellm.Router) -> None: + router.cooldown_cache.add_deployment_to_cooldown( + model_id="original", original_exception=RuntimeError("test outage"), exception_status=503, cooldown_time=3600 + ) + + +async def _record_replacement_success( + router: litellm.Router, deployment_id: str = "replacement", response: object = None +) -> None: + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + await router.deployment_callback_on_success( + {"litellm_params": _repin_params(deployment_id), "standard_logging_object": {"total_tokens": 0}}, + {"choices": [{"message": {"content": "ok"}}]} if response is None else response, + timestamp, + timestamp, + ) + + +@pytest.mark.parametrize("enabled", [None, False, True], ids=["default", "disabled", "enabled"]) +@pytest.mark.parametrize("unavailable", ["cooldown", "removed", "temporary-filter"]) +async def test_session_repin_requires_opt_in_unavailable_original_and_success( + enabled: bool | None, unavailable: str +) -> None: + router: Final = _repinning_router(enabled) + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + await router.cache.async_set_cache(_repin_key("caller-2"), {"model_id": "original"}) + user_key: Final = DeploymentAffinityCheck.get_affinity_cache_key("repin-group", "caller-1") + await router.cache.async_set_cache(user_key, {"model_id": "original"}) + if unavailable == "cooldown": + _cool_down_original(router) + elif unavailable == "removed": + router.delete_deployment("original") + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + await callback.async_pre_call_deployment_hook(_repin_params("replacement"), None) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + await _record_replacement_success(router) + expected: Final = "replacement" if enabled and unavailable != "temporary-filter" else "original" + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected} + assert await router.cache.async_get_cache(_repin_key("caller-2")) == {"model_id": "original"} + assert await router.cache.async_get_cache(user_key) == {"model_id": "original"} + router.cooldown_cache.cooldown_store.flush_cache() + await _record_replacement_success(router, "other") + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected} + finally: + router.discard() + + +@pytest.mark.parametrize("status", ["failed", "incomplete", "in_progress", "cancelled"]) +@pytest.mark.parametrize("wrapped", [False, True]) +async def test_unsuccessful_responses_do_not_repin(status: str, wrapped: bool) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + response: Final[Mapping[str, object]] = {"status": status} + await _record_replacement_success(router, response={"response": response} if wrapped else response) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + finally: + router.discard() + + +async def test_concurrent_successful_replacements_keep_first_winner() -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + await asyncio.gather(_record_replacement_success(router), _record_replacement_success(router, "other")) + winner: Final = await router.cache.async_get_cache(_repin_key()) + assert winner in ({"model_id": "replacement"}, {"model_id": "other"}) + router.cooldown_cache.cooldown_store.flush_cache() + await _record_replacement_success(router, "original") + assert await router.cache.async_get_cache(_repin_key()) == winner + finally: + router.discard() + + +async def test_healthy_legacy_string_pin_is_preserved_with_repinning_enabled() -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), "original") + await _record_replacement_success(router) + assert await router.cache.async_get_cache(_repin_key()) == "original" + finally: + router.discard() + + +@pytest.mark.parametrize("route_bucket", ["metadata", "litellm_metadata"]) +async def test_repin_session_after_success_reads_split_metadata(route_bucket: str) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + session_bucket: Final = "litellm_metadata" if route_bucket == "metadata" else "metadata" + await router._repin_session_after_success( + { + "model_info": {"id": "replacement"}, + route_bucket: {"deployment_model_name": "repin-group"}, + session_bucket: {"session_id": "repin-session", "user_api_key_hash": "caller-1"}, + }, + {"status": "completed"}, + ) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "replacement"} + finally: + router.discard() + + +@pytest.mark.parametrize("enabled", [False, True]) +async def test_session_selection_observes_another_workers_repin_only_when_enabled(enabled: bool) -> None: + server: Final = FakeServer() + sync_client: Final = FakeRedis(server=server) + pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server) + router: Final = _repinning_router(enabled) + try: + with ( + patch("redis.Redis.from_url", return_value=sync_client), + patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), + ): + redis: Final = RedisCache(url=f"redis://repin-cache.example:6379/{int(enabled)}") + router.cache.attach_redis_cache(redis) + other_worker: Final = DualCache(redis_cache=redis) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + await other_worker.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + callback: Final = next( + cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck) + ) + selected: Final = await callback.async_filter_deployments( + model="repin-group", + healthy_deployments=router.get_model_list(model_name="repin-group") or [], + messages=None, + request_kwargs=_repin_params("replacement"), + ) + assert [deployment["model_info"]["id"] for deployment in selected] == [ + "replacement" if enabled else "original" + ] + finally: + router.discard() + await pool.aclose() + sync_client.close() + + +@pytest.fixture +async def shared_session_router() -> AsyncIterator[tuple[litellm.Router, FakeServer, RedisCache]]: + server: Final = FakeServer() + sync_client: Final = FakeRedis(server=server) + pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server) + router: Final = _repinning_router() + try: + with ( + patch("redis.Redis.from_url", return_value=sync_client), + patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), + ): + redis: Final = RedisCache(url="redis://session-outage.example:6379/0") + router.cache.attach_redis_cache(redis) + yield router, server, redis + finally: + router.discard() + await pool.aclose() + sync_client.close() + + +async def _selected_session_ids(router: litellm.Router) -> tuple[str, ...]: + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + selected: Final = await callback.async_filter_deployments( + model="repin-group", + healthy_deployments=router.get_model_list(model_name="repin-group") or [], + messages=None, + request_kwargs=_repin_params("replacement"), + ) + return tuple(deployment["model_info"]["id"] for deployment in selected) + + +@pytest.mark.parametrize("has_pin", [False, True]) +async def test_redis_read_failures_and_open_breaker_preserve_session_routing( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], has_pin: bool +) -> None: + router, server, redis = shared_session_router + if has_pin: + await redis.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + assert await _selected_session_ids(router) == ("replacement",) + server.connected = False + with patch("time.time", return_value=1790000000.0): + for _ in range(redis._circuit_breaker.failure_threshold + 1): + assert set(await _selected_session_ids(router)) == ( + {"replacement"} if has_pin else {"original", "replacement", "other"} + ) + assert redis._circuit_breaker.is_open() + + +async def test_confirmed_redis_miss_does_not_revive_a_local_pin_during_outage( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], +) -> None: + router, server, _ = shared_session_router + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + assert set(await _selected_session_ids(router)) == {"original", "replacement", "other"} + server.connected = False + assert set(await _selected_session_ids(router)) == {"original", "replacement", "other"} + + +async def test_shared_session_pin_resumes_after_a_transient_redis_failure( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], +) -> None: + router, server, redis = shared_session_router + await redis.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + server.connected = False + assert await _selected_session_ids(router) == ("original",) + server.connected = True + assert await _selected_session_ids(router) == ("replacement",) + + +@pytest.mark.parametrize("shared_value", [None, {"model_id": "original"}]) +async def test_delayed_shared_pin_read_preserves_a_newer_local_write( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], shared_value: Mapping[str, str] | None +) -> None: + router, server, redis = shared_session_router + if shared_value is not None: + await redis.async_set_cache(_repin_key(), shared_value, ttl=3600) + observed: Final = asyncio.Event() + resume: Final = asyncio.Event() + + class DelayedReader(Redis): + async def get(self, name: str) -> object: + value: Final = await super().get(name) + observed.set() + await resume.wait() + return value + + reader: Final = DelayedReader(connection_pool=redis.init_async_client().connection_pool) + litellm.in_memory_llm_clients_cache.set_cache(key=redis._get_async_client_cache_key(), value=reader) + pending: Final = asyncio.create_task(_selected_session_ids(router)) + try: + await observed.wait() + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + finally: + resume.set() + await pending + server.connected = False + assert await _selected_session_ids(router) == ("replacement",) + + +@pytest.mark.parametrize( + "response", + [ + {"status": "failed", "response": {"status": "completed"}}, + {"error": {"message": "test failure"}, "response": {"status": "completed"}}, + {"response": {"response": {"status": "completed"}}}, + ], +) +async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: Mapping[str, object]) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + await _record_replacement_success(router, response=response) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + finally: + router.discard() + + +class _ReadCountingCache(InMemoryCache): + def __init__(self) -> None: + super().__init__() + self.read_keys: tuple[str, ...] = () + + def get_cache(self, key: str, **kwargs: object) -> object: + self.read_keys += (key,) + return super().get_cache(key, **kwargs) + + +@pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"]) +async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str) -> None: + router: Final = _repinning_router(case != "disabled") + reads: Final = _ReadCountingCache() + try: + router.cooldown_cache.cooldown_store.in_memory_cache = reads + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + if case == "group-disabled": + callback.model_group_affinity_config = {"repin-group": []} + await router._repin_session_after_success( + { + "model_info": {"id": "replacement"}, + "metadata": { + "deployment_model_name": "repin-group", + **({"session_id": "repin-session"} if case != "no-session" else {}), + SESSION_ID_GENERATED_METADATA_KEY: case == "generated-session", + }, + }, + {"status": "completed"}, + ) + expected: Final = ( + { + router.cooldown_cache.get_cooldown_cache_key(model_id) + for model_id in ("original", "replacement", "other") + } + if case == "active" + else set() + ) + assert set(reads.read_keys) == expected + finally: + router.discard() + + class MockResponse: def __init__(self, json_data, status_code): self._json_data = json_data diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 09ce002353f..75b6ab04446 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -2,7 +2,7 @@ import asyncio import importlib import json import time -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from datetime import datetime, timedelta from types import MappingProxyType from typing import Any, Final, Literal, NoReturn @@ -11,6 +11,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest import respx +from pydantic import TypeAdapter import litellm from litellm import Router @@ -19,7 +20,8 @@ from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import Fa from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure +from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments, mark_advisor_orchestration_failure +from litellm.types.utils import Choices, Message, ModelResponse, StandardLoggingPayload, Usage from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, PRE_ROUTING_SELECTED_MODEL_KEY, @@ -147,6 +149,220 @@ class RecordingRouter: return StreamingWrapper() +def _request_logging_obj() -> Logging: + return Logging( + model="primary-model", + messages=[{"role": "user", "content": "ping"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1), + litellm_call_id="fallback-hop", + function_id="fallback-hop", + ) + + +class _SuccessRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: tuple[tuple[int, float], ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + self.logged = (*self.logged, (payload["completion_tokens"], payload["response_cost"])) + + +class HopLoggingItsOwnSuccessRouter(FakeRouter): + async def async_function_with_fallbacks(self, *args: object, **kwargs: object) -> StreamingWrapper: + hop_response: Final = ModelResponse( + id="chatcmpl-hop", + model="gpt-5.6", + choices=[Choices(index=0, message=Message(role="assistant", content="served by the hop"))], + usage=Usage(prompt_tokens=52, completion_tokens=7, total_tokens=59), + ) + now: Final = datetime(2026, 1, 1) + logging_obj: Final = kwargs["litellm_logging_obj"] + assert isinstance(logging_obj, Logging) + await logging_obj.async_success_handler(result=hop_response, start_time=now, end_time=now) + return StreamingWrapper() + + +@pytest.mark.asyncio +async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_partial_usage() -> None: + """A stream the failed attempt dropped stashed the usage it consumed on the request's logging object. The + hop that replaces it logs its own usage, so the integrations see the hop's tokens and cost, not the stash.""" + recorder: Final = _SuccessRecorder() + logging_obj: Final = Logging( + model="gpt-5.6", + messages=[{"role": "user", "content": "ping"}], + stream=False, + call_type="completion", + start_time=datetime(2026, 1, 1), + litellm_call_id="fallback-hop-success", + function_id="fallback-hop-success", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + logging_obj.model_call_details["litellm_params"] = {"acompletion": True} + logging_obj.model_call_details["current_attempt_api_call_start_time"] = datetime(2026, 1, 1) + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) + + await run_async_fallback( + litellm_router=HopLoggingItsOwnSuccessRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + ) + + assert tuple(tokens for tokens, _ in recorder.logged) == (7,) + assert recorder.logged[0][1] != 0.01 + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + assert logging_obj.model_call_details["current_attempt_api_call_start_time"] is None + + +@pytest.mark.asyncio +async def test_run_async_fallback_preserves_superseded_usage_when_hop_fails_before_opening() -> None: + logging_obj: Final = _request_logging_obj() + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=AlwaysFailRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + ) + + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + + +class _RateLimitedHop(Exception): + def __init__(self, deployment_id: str = "fallback-dep-1") -> None: + super().__init__("fallback deployment rate limited") + self.status_code = 429 + self.failed_deployment_id = deployment_id + + +_EARLIER_FAILURE_MARKER: Final = datetime(2026, 1, 1) +_FIRST_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 2) +_FINAL_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 3) + + +class _TwoFailureFallbackRouter(Router): + def __init__(self) -> None: + super().__init__( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"}, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + allowed_fails=0, + cooldown_time=60, + num_retries=0, + ) + + async def async_function_with_fallbacks(self, *args: object, **kwargs: object) -> NoReturn: + logging_obj: Final[object] = kwargs["litellm_logging_obj"] + assert isinstance(logging_obj, Logging) + assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False + + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FIRST_HOP_FAILURE_MARKER + try: + raise _RateLimitedHop("fallback-dep-1") + except _RateLimitedHop: + logging_obj.mark_logging_complete("async_failure") + + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER + raise _RateLimitedHop("fallback-dep-2") + + +def _two_deployment_router() -> Router: + return Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"}, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + num_retries=0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_already_logged", [False, True]) +async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool) -> None: + """The manual cooldown stands in for the failure callbacks the logging object blocks once a failure has + been logged. A hop whose own failure was the first one logged already ran them, so only a hop that failed + with the flag already set gets the manual cooldown.""" + router: Final = _two_deployment_router() + logging_obj: Final = _request_logging_obj() + if failure_already_logged: + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _EARLIER_FAILURE_MARKER + logging_obj.mark_logging_complete("async_failure") + + async def rate_limited_hop(*args: object, **kwargs: object) -> NoReturn: + hop_logging_obj: Final = kwargs["litellm_logging_obj"] + assert isinstance(hop_logging_obj, Logging) + hop_logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER + if hop_logging_obj.model_call_details.get("has_logged_async_failure", False) is not True: + hop_logging_obj.mark_logging_complete("async_failure") + raise _RateLimitedHop() + + with pytest.raises(_RateLimitedHop): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + original_function=rate_limited_hop, + metadata={}, + ) + + cooled_down: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + assert cooled_down == (["fallback-dep-1"] if failure_already_logged else []) + + +@pytest.mark.asyncio +async def test_run_async_fallback_cools_down_the_final_deployment_after_retry_failure_logging() -> None: + router: Final = _TwoFailureFallbackRouter() + logging_obj: Final = _request_logging_obj() + assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False + + with pytest.raises(_RateLimitedHop): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + metadata={}, + ) + + assert logging_obj.model_call_details.get("has_logged_async_failure", False) is True + cooled_down: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + assert cooled_down == ["fallback-dep-2"] + + @pytest.mark.asyncio async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_call(): """A nested fallback (multi-hop) must keep collecting errors, so the opt-in @@ -1163,11 +1379,18 @@ class TestRunAsyncFallbackTriggersCooldown: return None async def async_function_with_fallbacks(self, *args, **kwargs): + logging_obj: Final = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER + if logging_obj.model_call_details.get("has_logged_async_failure", False) is not True: + logging_obj.mark_logging_complete("async_failure") raise RuntimeError("fallback model also failed") - def _logging_obj(self, has_logged_async_failure: bool) -> MagicMock: - logging_obj = MagicMock() - logging_obj.model_call_details = {"has_logged_async_failure": has_logged_async_failure} + def _logging_obj(self, has_logged_async_failure: bool) -> Logging: + logging_obj: Final = _request_logging_obj() + if has_logged_async_failure: + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _EARLIER_FAILURE_MARKER + logging_obj.mark_logging_complete("async_failure") return logging_obj @pytest.mark.asyncio diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 726128386d5..78722458cf3 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -116,6 +116,9 @@ def _wrap_nested_level(value: object, depth: int) -> object: (responses_detectors.provider_affinity, {"provider_affinity_header": ""}, True), (responses_detectors.provider_affinity, {"provider_affinity_header": None, "session_id": "session-1"}, False), (responses_detectors.provider_affinity, {}, False), + (responses_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/openai/v1"}, True), + (messages_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/anthropic/v1"}, True), + (messages_detectors.provider_affinity, {"api_base": "https://other.example/v1"}, False), (detectors.streaming, {"stream": True}, True), (detectors.streaming, {"stream": False}, False), (detectors.timeout_variants, {"request_timeout": 5}, True), diff --git a/tests/unit/secret_managers/test_google_secret_manager.py b/tests/unit/secret_managers/test_google_secret_manager.py index 2051be08f7c..7effde4fb30 100644 --- a/tests/unit/secret_managers/test_google_secret_manager.py +++ b/tests/unit/secret_managers/test_google_secret_manager.py @@ -1,6 +1,7 @@ import base64 import binascii import json +from collections import OrderedDict from dataclasses import dataclass from typing import Final @@ -120,7 +121,7 @@ def manager(monkeypatch: pytest.MonkeyPatch) -> GoogleSecretManager: monkeypatch.setattr(litellm, "secret_manager_client", None) monkeypatch.setattr(litellm, "_key_management_system", None) monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings()) - monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", {}) + monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", OrderedDict()) monkeypatch.setattr(google.auth, "default", _default_credentials) monkeypatch.setenv("GOOGLE_SECRET_MANAGER_PROJECT_ID", PROJECT_ID) monkeypatch.delenv("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", raising=False) diff --git a/tests/unit/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py index 2a85122d91c..deba7fe9587 100644 --- a/tests/unit/test_anthropic_beta_headers_filtering.py +++ b/tests/unit/test_anthropic_beta_headers_filtering.py @@ -17,7 +17,9 @@ import pytest import litellm from litellm.anthropic_beta_headers_manager import ( + GetAnthropicBetaHeadersConfig, filter_and_transform_beta_headers, + get_beta_headers_config, update_headers_with_filtered_beta, update_request_with_filtered_beta, ) @@ -429,13 +431,13 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["fine-grained-tool-streaming-2025-05-14"] @pytest.mark.parametrize( - "provider", ["anthropic", "bedrock", "bedrock_converse", "vertex_ai", "databricks"] + "provider", ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai", "databricks"] ) def test_thinking_binding_controls_forwarded(self, provider): """`thinking.block_binding` (preserved thinking, Claude Fable 5.1) is only accepted alongside thinking-binding-controls-2026-08-01. The body field is forwarded untouched, so stripping the header (previously unknown, hence - dropped) makes Bedrock and Vertex reject the request with + dropped) makes Bedrock, Vertex and Azure AI Foundry reject the request with "thinking.adaptive.block_binding: Extra inputs are not permitted".""" filtered = filter_and_transform_beta_headers( beta_headers=["thinking-binding-controls-2026-08-01"], @@ -558,3 +560,13 @@ class TestAnthropicBetaHeadersFiltering: headers = {"anthropic-version": "2023-06-01"} assert update_headers_with_filtered_beta(headers, "bedrock_mantle") == {"anthropic-version": "2023-06-01"} + + +@pytest.mark.parametrize("flag", ["LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "LITELLM_OFFLINE_MODE"]) +def test_local_flags_load_bundled_beta_headers_without_fetching(monkeypatch, flag): + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + monkeypatch.setenv(flag, "True") + with patch("litellm.anthropic_beta_headers_manager.httpx.get") as remote_get: + config = get_beta_headers_config(url="https://example.test/anthropic_beta_headers_config.json") + remote_get.assert_not_called() + assert config == GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() diff --git a/tests/unit/test_auto_update_price_and_context_window_file.py b/tests/unit/test_auto_update_price_and_context_window_file.py deleted file mode 100644 index 435747e9a09..00000000000 --- a/tests/unit/test_auto_update_price_and_context_window_file.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Unit tests for the Friendli transform in -`.github/scripts/auto_update_price_and_context_window_file.py`.""" - -from __future__ import annotations - -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] - / ".github" - / "scripts" - / "auto_update_price_and_context_window_file.py" -) - - -@pytest.fixture(scope="module") -def sync_module(): - spec = importlib.util.spec_from_file_location( - "auto_update_price_and_context_window_file", SCRIPT_PATH - ) - assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" - module = importlib.util.module_from_spec(spec) - sys.modules["auto_update_price_and_context_window_file"] = module - spec.loader.exec_module(module) - return module - - -def _reasoning_model(**overrides: object) -> dict: - model = { - "id": "zai-org/GLM-Test", - "base_model": "zhipuai/glm-test", - "context_length": 1048576, - "max_completion_tokens": 131072, - "pricing": {"input": "0.00000015", "output": "0.0000005", "input_cache_read": "0.00000003"}, - "reasoning": True, - "reasoning_options": [{"type": "effort", "values": ["max", "high", "low"]}], - "functionality": { - "tool_call": True, - "parallel_tool_call": True, - "structured_output": True, - "system_messages": True, - "tool_choice": True, - }, - "input_modalities": ["text", "image", "video"], - "mode": "chat", - } - model.update(overrides) - return model - - -def test_transform_emits_declared_effort_levels_in_canonical_order(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_reasoning"] is True - assert entry["reasoning_effort_levels"] == ["low", "high", "max"] - assert not any(k.endswith("_reasoning_effort") for k in entry) - - -def test_transform_reasoning_model_without_effort_options_declares_empty_levels(sync_module): - model = _reasoning_model(reasoning_options=[{"type": "budget_tokens", "values": []}]) - entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"] - assert entry["reasoning_effort_levels"] == [] - - -def test_transform_non_reasoning_model_declares_no_levels(sync_module): - model = _reasoning_model(reasoning=False, reasoning_options=[]) - entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"] - assert entry["supports_reasoning"] is False - assert "reasoning_effort_levels" not in entry - - -def test_transform_max_tokens_mirrors_output_cap_not_context(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["max_input_tokens"] == 1048576 - assert entry["max_output_tokens"] == 131072 - assert entry["max_tokens"] == entry["max_output_tokens"] - - -def test_transform_prompt_caching_follows_cache_pricing(sync_module): - cached = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert cached["supports_prompt_caching"] is True - assert cached["cache_read_input_token_cost"] == 3e-08 - - uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"}) - uncached = sync_module.transform_friendli_data([uncached_model], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert uncached["supports_prompt_caching"] is False - assert "cache_read_input_token_cost" not in uncached - - -def test_transform_modalities_set_vision_image_and_video_flags(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_vision"] is True - assert entry["supports_image_input"] is True - assert entry["supports_video_input"] is True - - text_only = _reasoning_model(input_modalities=["text"]) - entry_text = sync_module.transform_friendli_data([text_only], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry_text["supports_vision"] is False - assert entry_text["supports_image_input"] is False - assert entry_text["supports_video_input"] is False - - -def test_transform_skips_rows_without_valid_token_prices_so_priced_local_entries_survive(sync_module): - local = { - "friendliai/zai-org/GLM-Test": { - "litellm_provider": "friendliai", - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - } - } - unpriced_rows = [ - _reasoning_model(pricing={}), - _reasoning_model(pricing=None), - _reasoning_model(pricing={"input": "0.00000015"}), - _reasoning_model(pricing={"output": "0.0000005"}), - _reasoning_model(pricing={"input": "not-a-number", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "-0.00000015", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "inf", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "nan", "output": "0.0000005"}), - ] - remote = sync_module.transform_friendli_data(unpriced_rows, local) - assert remote == {} - sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote)) - assert local["friendliai/zai-org/GLM-Test"]["input_cost_per_token"] == 1.5e-07 - assert local["friendliai/zai-org/GLM-Test"]["output_cost_per_token"] == 5e-07 - - -def test_transform_keeps_zero_priced_rows(sync_module): - free_model = _reasoning_model(pricing={"input": "0", "output": "0"}) - entry = sync_module.transform_friendli_data([free_model], {})["friendliai/zai-org/GLM-Test"] - assert entry["input_cost_per_token"] == 0.0 - assert entry["output_cost_per_token"] == 0.0 - - -def test_transforms_survive_failed_fetch(sync_module): - assert sync_module.transform_friendli_data(None, {}) == {} - assert sync_module.transform_friendli_data([], {}) == {} - assert sync_module.transform_openrouter_data(None) == {} - assert sync_module.transform_vercel_ai_gateway_data(None) == {} - - -def test_vercel_transform_skips_rows_without_token_pricing_or_limits(sync_module): - rows = [ - { - "id": "wan-video", - "pricing": {"video_duration_pricing": [{"resolution": "720p", "cost_per_second": "0.1"}]}, - }, - { - "id": "qwen3-embedding", - "context_window": 32768, - "max_tokens": 32768, - "pricing": {"input": "0.00000001"}, - }, - { - "id": "no-limits-chat", - "pricing": {"input": "0.000001", "output": "0.000002"}, - }, - { - "id": "good-chat", - "context_window": 128000, - "max_tokens": 8192, - "pricing": {"input": "0.000001", "output": "0.000002"}, - }, - ] - transformed = sync_module.transform_vercel_ai_gateway_data(rows) - assert list(transformed) == ["vercel_ai_gateway/good-chat"] - assert transformed["vercel_ai_gateway/good-chat"]["input_cost_per_token"] == 1e-06 - assert transformed["vercel_ai_gateway/good-chat"]["output_cost_per_token"] == 2e-06 - - -def test_sync_replaces_friendli_entries_so_dropped_cache_pricing_does_not_survive(sync_module): - local = { - "friendliai/zai-org/GLM-Test": { - "litellm_provider": "friendliai", - "cache_read_input_token_cost": 3e-08, - "supports_prompt_caching": True, - } - } - uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"}) - remote = sync_module.transform_friendli_data([uncached_model], local) - sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote)) - synced = local["friendliai/zai-org/GLM-Test"] - assert "cache_read_input_token_cost" not in synced - assert synced["supports_prompt_caching"] is False - - -def test_sync_still_merges_entries_outside_replace_keys(sync_module): - local = {"openrouter/some-model": {"input_cost_per_token": 1e-06, "supports_vision": True}} - remote = {"openrouter/some-model": {"input_cost_per_token": 2e-06}} - sync_module.sync_local_data_with_remote(local, remote) - assert local["openrouter/some-model"] == {"input_cost_per_token": 2e-06, "supports_vision": True} - - -def test_transform_inherits_allowlisted_keys_from_base_model_entry(sync_module): - local = { - "zhipuai/glm-test": { - "supports_pdf_input": True, - "supports_assistant_prefill": True, - "input_cost_per_token": 9e-06, - } - } - entry = sync_module.transform_friendli_data([_reasoning_model()], local)[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_pdf_input"] is True - assert entry["supports_assistant_prefill"] is True - assert entry["input_cost_per_token"] == 1.5e-07 diff --git a/tests/unit/test_cerebras_llama_model_metadata.py b/tests/unit/test_cerebras_llama_model_metadata.py new file mode 100644 index 00000000000..c5ec096e553 --- /dev/null +++ b/tests/unit/test_cerebras_llama_model_metadata.py @@ -0,0 +1,115 @@ +# What does this test? +""" +Verifies that Cerebras Llama 3.1 models (standard hyphenated aliases): +- cerebras/llama-3.1-8b +- cerebras/llama-3.1-70b + +have correct rate card metadata, context window limits, and cost calculation parity +with the un-hyphenated variants in model_prices_and_context_window.json. +Fixes: https://github.com/BerriAI/litellm/issues/48157 +""" + +import json +from pathlib import Path + +import pytest + +import litellm + + +@pytest.fixture(scope="module") +def model_cost_maps(): + root = Path(__file__).resolve().parents[2] + main_path = root / "model_prices_and_context_window.json" + backup_path = root / "litellm" / "model_prices_and_context_window_backup.json" + + with open(main_path, "r", encoding="utf-8") as f: + main_map = json.load(f) + with open(backup_path, "r", encoding="utf-8") as f: + backup_map = json.load(f) + + return main_map, backup_map + + +CEREBRAS_MODELS_EXPECTED = { + "cerebras/llama-3.1-8b": { + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1e-07, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "mode": "chat", + "litellm_provider": "cerebras", + }, + "cerebras/llama-3.1-70b": { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "mode": "chat", + "litellm_provider": "cerebras", + }, +} + +PAIR_PARITY = [ + ("cerebras/llama-3.1-8b", "cerebras/llama3.1-8b"), + ("cerebras/llama-3.1-70b", "cerebras/llama3.1-70b"), +] + + +@pytest.mark.parametrize("model_id,expected_data", CEREBRAS_MODELS_EXPECTED.items()) +def test_cerebras_llama_model_metadata(model_cost_maps, model_id, expected_data): + """Test that hyphenated Cerebras Llama 3.1 models exist and have expected metadata in both catalogs.""" + main_map, backup_map = model_cost_maps + + assert model_id in main_map, f"Missing {model_id} in main catalog" + assert model_id in backup_map, f"Missing {model_id} in backup catalog" + + main_entry = main_map[model_id] + backup_entry = backup_map[model_id] + assert main_entry == backup_entry, f"Parity mismatch between main and backup for {model_id}" + + for key, expected_val in expected_data.items(): + assert main_entry.get(key) == expected_val, ( + f"{model_id}[{key}] got {main_entry.get(key)}, expected {expected_val}" + ) + + +@pytest.mark.parametrize("hyphenated,unhyphenated", PAIR_PARITY) +def test_cerebras_llama_parity_with_unhyphenated(model_cost_maps, hyphenated, unhyphenated): + """Test that hyphenated and un-hyphenated variants have identical pricing and context limits.""" + main_map, _ = model_cost_maps + + assert hyphenated in main_map + assert unhyphenated in main_map + + entry_hyphen = main_map[hyphenated] + entry_unhyphen = main_map[unhyphenated] + + assert entry_hyphen["input_cost_per_token"] == entry_unhyphen["input_cost_per_token"] + assert entry_hyphen["output_cost_per_token"] == entry_unhyphen["output_cost_per_token"] + assert entry_hyphen["max_input_tokens"] == entry_unhyphen["max_input_tokens"] + assert entry_hyphen["max_output_tokens"] == entry_unhyphen["max_output_tokens"] + + +def test_cerebras_llama_cost_per_token(): + """Test that litellm.cost_per_token accurately computes costs for hyphenated Cerebras models.""" + prompt_tokens = 1_000_000 + completion_tokens = 1_000_000 + + # llama-3.1-8b: $0.10 input + $0.10 output per 1M tokens + cost_prompt, cost_comp = litellm.cost_per_token( + model="cerebras/llama-3.1-8b", + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + assert pytest.approx(cost_prompt, rel=1e-5) == 0.10 + assert pytest.approx(cost_comp, rel=1e-5) == 0.10 + + # llama-3.1-70b: $0.60 input + $0.60 output per 1M tokens + cost_prompt_70b, cost_comp_70b = litellm.cost_per_token( + model="cerebras/llama-3.1-70b", + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + assert pytest.approx(cost_prompt_70b, rel=1e-5) == 0.60 + assert pytest.approx(cost_comp_70b, rel=1e-5) == 0.60 diff --git a/tests/unit/test_circleci_path_filter.py b/tests/unit/test_circleci_path_filter.py index f94b2fa863a..c44560b96f1 100644 --- a/tests/unit/test_circleci_path_filter.py +++ b/tests/unit/test_circleci_path_filter.py @@ -96,7 +96,6 @@ CI = [".github/workflows/test-unit.yml"] ("ui", BACKEND + CLIENT, "run"), ("cost-map-only", ["model_prices_and_context_window.json"], "run"), ("cost-map-only", ["litellm/model_prices_and_context_window_backup.json"], "run"), - ("cost-map-only", ["model_prices_and_context_window.schema.json"], "run"), ( "cost-map-only", ["model_prices_and_context_window.json", "tests/unit/proxy/test_y.py"], diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 2854f9010b2..27b25c1f98e 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -16,6 +16,7 @@ import litellm 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.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.fallback_generalizations import ( get_fallback_generalization_rules, @@ -31,10 +32,10 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.fireworks_ai.cost_calculator import get_base_model_for_pricing from litellm.llms.together_ai.cost_calculator import get_model_params_and_category from litellm.types.llms.base import CachedTokensDetails -from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.llms.openai import HttpxBinaryResponseContent, OpenAIRealtimeStreamList, ResponseAPIUsage, ResponsesAPIResponse from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.rerank import RerankResponse -from litellm.types.utils import CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage +from litellm.types.utils import CacheCreationTokenDetails, CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject from litellm.utils import supports_prompt_caching import os @@ -3337,14 +3338,7 @@ def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models(_loca def test_gemini_cache_tokens_details_no_negative_values(): - """ - Test for Issue #18750: Negative text_tokens with Gemini caching - - When using Gemini with explicit caching, the response includes cacheTokensDetails - which breaks down cached tokens by modality. This test ensures that: - 1. text_tokens is never negative - 2. We correctly subtract cached tokens per modality (not total) - """ + """Ensure issue #18750 cache details remain non-negative after cost parsing.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -3372,27 +3366,17 @@ def test_gemini_cache_tokens_details_no_negative_values(): usage = VertexGeminiConfig.calculate_usage(completion_response) - # Text tokens should be non-cached text only: 9402 - 9393 = 9 - assert usage.prompt_tokens_details.text_tokens == 9, ( - f"Expected text_tokens=9, got {usage.prompt_tokens_details.text_tokens}" - ) + prompt_details: Final = parse_prompt_tokens_details(usage) - # Image tokens should be non-cached image only: 258 - 258 = 0 - assert usage.prompt_tokens_details.image_tokens == 0, ( - f"Expected image_tokens=0, got {usage.prompt_tokens_details.image_tokens}" + assert usage.prompt_tokens_details.text_tokens == 9402 + assert usage.prompt_tokens_details.image_tokens == 258 + assert usage.prompt_tokens_details.cached_tokens == 9651 + assert usage.prompt_tokens_details.cached_tokens_details == CachedTokensDetails( + text_tokens=9393, image_tokens=258 ) - - # Total cached should match - assert usage.prompt_tokens_details.cached_tokens == 9651, ( - f"Expected cached_tokens=9651, got {usage.prompt_tokens_details.cached_tokens}" - ) - - # MOST IMPORTANT: text_tokens should NEVER be negative - assert usage.prompt_tokens_details.text_tokens >= 0, ( - f"BUG: text_tokens is negative ({usage.prompt_tokens_details.text_tokens})! This was the issue in #18750" - ) - - print("✅ Issue #18750 fix verified: text_tokens is correctly calculated and non-negative") + assert prompt_details["text_tokens"] == 9 + assert prompt_details["image_tokens"] == 0 + assert prompt_details["text_tokens"] >= 0 def test_gemini_without_cache_tokens_details(): @@ -5963,6 +5947,110 @@ def test_batch_cost_calculator_bills_cache_write_tokens_at_the_batch_input_rate_ assert prompt_cost == pytest.approx(300_048 * 3e-6) +_TWELVE_HOUR_BATCH_MODEL_INFO: Final = cast( + ModelInfo, + { + "input_cost_per_token_batches": 1e-6, + "input_cost_per_token_above_272k_tokens_batches": 3e-6, + "cache_creation_input_token_cost_batches": 1.25e-6, + "cache_creation_input_token_cost_above_272k_tokens_batches": 3.75e-6, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 7.5e-6, + "output_cost_per_token_batches": 4e-6, + }, +) + + +def _twelve_hour_cache_write_usage(prompt_tokens: int) -> Usage: + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=4, + total_tokens=prompt_tokens + 4, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=0, + cache_creation_tokens=500, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_30m_input_tokens=300, ephemeral_12h_input_tokens=200 + ), + ), + ) + + +@pytest.mark.parametrize( + ("prompt_tokens", "input_rate", "write_rate", "twelve_hour_rate"), + [(1_000, 1e-6, 1.25e-6, 2.5e-6), (272_000, 1e-6, 1.25e-6, 2.5e-6), (300_000, 3e-6, 3.75e-6, 7.5e-6)], +) +def test_batch_cost_calculator_bills_12h_cache_writes_at_the_12h_batch_rate_for_the_crossed_tier( + prompt_tokens: int, input_rate: float, write_rate: float, twelve_hour_rate: float +) -> None: + from litellm.cost_calculator import batch_cost_calculator + + prompt_cost, completion_cost_value = batch_cost_calculator( + usage=_twelve_hour_cache_write_usage(prompt_tokens), + model="lit-12h-batch-model", + custom_llm_provider="openai", + model_info=_TWELVE_HOUR_BATCH_MODEL_INFO, + ) + + assert prompt_cost == pytest.approx((prompt_tokens - 500) * input_rate + 300 * write_rate + 200 * twelve_hour_rate) + assert completion_cost_value == pytest.approx(4 * 4e-6) + + +@pytest.mark.parametrize( + ("prompt_tokens", "input_rate", "twelve_hour_write_rate"), [(1_000, 1e-6, 2.5e-6), (300_000, 3e-6, 7.5e-6)] +) +def test_batch_cost_calculator_bills_12h_cache_writes_once_when_only_the_12h_batch_key_exists( + prompt_tokens: int, input_rate: float, twelve_hour_write_rate: float +) -> None: + from litellm.cost_calculator import batch_cost_calculator + + model_info: Final = cast( + ModelInfo, + { + key: rate + for key, rate in _TWELVE_HOUR_BATCH_MODEL_INFO.items() + if not key.startswith("cache_creation_input_token_cost_") + or key.startswith("cache_creation_input_token_cost_above_12hr") + }, + ) + + prompt_cost, _ = batch_cost_calculator( + usage=_twelve_hour_cache_write_usage(prompt_tokens), + model="lit-12h-batch-model", + custom_llm_provider="openai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx((prompt_tokens - 200) * input_rate + 200 * twelve_hour_write_rate) + + +@pytest.mark.parametrize(("prompt_tokens", "write_rate"), [(1_000, 1.25e-6), (300_000, 3.75e-6)]) +def test_batch_cost_calculator_bills_12h_cache_writes_at_the_batch_cache_write_rate_without_a_12h_batch_key( + prompt_tokens: int, write_rate: float +) -> None: + from litellm.cost_calculator import batch_cost_calculator + + model_info: Final = cast( + ModelInfo, + { + key: rate + for key, rate in _TWELVE_HOUR_BATCH_MODEL_INFO.items() + if not key.startswith("cache_creation_input_token_cost_above_12hr") + }, + ) + + prompt_cost, _ = batch_cost_calculator( + usage=_twelve_hour_cache_write_usage(prompt_tokens), + model="lit-12h-batch-model", + custom_llm_provider="openai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx( + (prompt_tokens - 500) * (1e-6 if prompt_tokens <= 272_000 else 3e-6) + 500 * write_rate + ) + + def test_batch_cost_calculator_prices_modalities_and_cached_tokens_together_in_the_crossed_tier() -> None: from litellm.cost_calculator import batch_cost_calculator @@ -5990,6 +6078,45 @@ def test_batch_cost_calculator_prices_modalities_and_cached_tokens_together_in_t assert prompt_cost == pytest.approx(298_926 * 3e-6 + 64 * 5e-6 + 10 * 3e-6 + 1_000 * 3e-7) +def test_batch_cost_calculator_prices_cached_audio_at_its_batch_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_token_batches": 1e-6, + "input_cost_per_audio_token": 6e-6, + "input_cost_per_audio_token_batches": 3e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_token_cost_batches": 4e-7, + "cache_read_input_audio_token_cost": 5e-7, + "cache_read_input_audio_token_cost_batches": 2e-7, + } + usage: Final = Usage( + prompt_tokens=40, + completion_tokens=0, + total_tokens=40, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, + audio_tokens=20, + cached_tokens=10, + cached_tokens_details=CachedTokensDetails(text_tokens=5, audio_tokens=5), + ), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-audio-batch-test-model", + custom_llm_provider="gemini", + model_info=model_info, + ) + + expected: Final = ( + 15 * model_info["input_cost_per_token_batches"] + + 15 * model_info["input_cost_per_audio_token_batches"] + + 5 * model_info["cache_read_input_token_cost_batches"] + + 5 * model_info["cache_read_input_audio_token_cost_batches"] + ) + assert prompt_cost == pytest.approx(expected) + + QWEN3_NEXT_REGIONS: Final = ("ap-northeast-1", "ap-south-1", "ap-southeast-2", "eu-west-1", "eu-west-2", "sa-east-1") @@ -9134,6 +9261,33 @@ def test_azure_tts_cost_uses_test_owned_character_rate(monkeypatch: pytest.Monke assert cost == pytest.approx(0.05) +def test_speech_cost_with_custom_pricing_bills_the_reported_usage() -> None: + usage: Final = Usage( + prompt_tokens=10, + completion_tokens=1000, + total_tokens=1010, + completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=1000), + ) + response: Final = HttpxBinaryResponseContent( + httpx.Response(200, content=b"RIFF....WAVE", headers={"content-type": "audio/wav"}) + ) + response.usage = usage + rates: Final = MappingProxyType({"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-5}) + + cost: Final = completion_cost( + completion_response=response, + model="vertex_ai/gemini-3.1-flash-tts-preview", + call_type="speech", + prompt="Say hello", + custom_cost_per_token=dict(rates), + ) + + assert cost == pytest.approx( + usage.prompt_tokens * rates["input_cost_per_token"] + + usage.completion_tokens * rates["output_cost_per_token"] + ) + + def test_moderation_cost_uses_test_owned_token_rate() -> None: prompt: Final = "moderation input" rates: Final = MappingProxyType( diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index a705597952b..578d5482822 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -186,13 +186,6 @@ 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) == () diff --git a/tests/unit/test_elevenlabs_tts_pricing.py b/tests/unit/test_elevenlabs_tts_pricing.py new file mode 100644 index 00000000000..16f726b9d6c --- /dev/null +++ b/tests/unit/test_elevenlabs_tts_pricing.py @@ -0,0 +1,49 @@ +"""ElevenLabs text-to-speech prices match the published API rates. + +https://elevenlabs.io/pricing/api lists "v3" and "v2 Multilingual" at $0.08 and +"Flash / Turbo" at $0.04, each "Price per 1K characters", and says API usage is +billed in US dollars, not credits. The map had v3 and Multilingual v2 at a +Scale-plan credit estimate ($0.18 per 1K characters) and no Flash rows. +""" + +import json +from pathlib import Path + +import pytest + +import litellm + +REPO_ROOT = Path(__file__).parents[2] +PRICES = { + "elevenlabs/eleven_v3": 8e-05, + "elevenlabs/eleven_multilingual_v2": 8e-05, + "elevenlabs/eleven_flash_v2_5": 4e-05, + "elevenlabs/eleven_flash_v2": 4e-05, +} + + +@pytest.fixture(autouse=True) +def local_model_cost_map(monkeypatch): + """Never depend on the published cost map having these rows yet.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model, per_character", PRICES.items()) +def test_metadata_and_backup(model, per_character): + main = json.loads((REPO_ROOT / "model_prices_and_context_window.json").read_text()) + backup = json.loads((REPO_ROOT / "litellm/model_prices_and_context_window_backup.json").read_text()) + assert main[model] == backup[model] + info = litellm.get_model_info(model) + assert info["litellm_provider"] == "elevenlabs" + assert info["mode"] == "audio_speech" + assert info["input_cost_per_character"] == per_character + + +@pytest.mark.parametrize("model, per_character", PRICES.items()) +def test_speech_cost_is_charged_per_character(model, per_character): + cost = litellm.completion_cost(model=model, call_type="speech", prompt="x" * 1000) + assert cost == pytest.approx(per_character * 1000) diff --git a/tests/unit/test_get_blog_posts.py b/tests/unit/test_get_blog_posts.py index 499dd9bcf42..d1de75a4500 100644 --- a/tests/unit/test_get_blog_posts.py +++ b/tests/unit/test_get_blog_posts.py @@ -190,8 +190,9 @@ def test_get_blog_posts_ttl_expired_refetches(): assert len(posts) == 2 -def test_get_blog_posts_local_env_var_skips_remote(monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_BLOG_POSTS", "true") +@pytest.mark.parametrize("flag", ["LITELLM_LOCAL_BLOG_POSTS", "LITELLM_OFFLINE_MODE"]) +def test_get_blog_posts_local_env_var_skips_remote(monkeypatch, flag): + monkeypatch.setenv(flag, "true") with patch("litellm.litellm_core_utils.get_blog_posts.httpx.get") as mock_get: posts = get_blog_posts(url=litellm.blog_posts_url) mock_get.assert_not_called() diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f5d13aecb6f..70b268bf2a2 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -20,7 +20,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import APITimeoutError, AsyncOpenAI +from openai import APITimeoutError, AsyncOpenAI, OpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion from pydantic import TypeAdapter @@ -95,6 +95,142 @@ def preserve_litellm_completion_state(monkeypatch: pytest.MonkeyPatch) -> None: WHITE_PNG: Final = (Path(__file__).parents[1] / "white_100x100.png").read_bytes() +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("client_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/zen/go/v1", None, "conversation-123"), + ("https://opencode.ai/zen/go/v1", "caller-override", "caller-override"), + ("https://other.example/v1", None, None), + ], +) +@pytest.mark.asyncio +async def test_completion_sends_session_only_to_the_actual_opencode_destination( + asynchronous: bool, + stream: bool, + client_base: str, + explicit_header: str | None, + expected_session: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENAI_BASE_URL", "https://opencode.ai/zen/go/v1") + monkeypatch.setenv("OPENAI_API_BASE", "https://opencode.ai/zen/go/v1") + monkeypatch.setattr(litellm, "api_base", "https://opencode.ai/zen/go/v1") + requests: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + + def respond(request: httpx.Request) -> httpx.Response: + requests.put_nowait(request) + payload: Final = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk" if stream else "chat.completion", + "created": 1, + "model": "session-test", + "choices": [ + { + "index": 0, + "delta" if stream else "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + } + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(payload)}\n\ndata: [DONE]\n\n", + ) + if stream + else httpx.Response(200, json=payload) + ) + + args: Final = { + "model": "openai/session-test", + "api_base": "https://opencode.ai/zen/go/v1", + "messages": [{"role": "user", "content": "hello"}], + "litellm_session_id": "conversation-123", + "extra_headers": {"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + "stream": stream, + } + if asynchronous: + async with AsyncOpenAI( + api_key="test-key", + base_url=client_base, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) as async_client: + response: Final = await acompletion(**args, client=async_client) + content: Final = ( + "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + if stream + else response.choices[0].message.content + ) + assert content == "OK" + else: + with OpenAI( + api_key="test-key", base_url=client_base, http_client=httpx.Client(transport=httpx.MockTransport(respond)) + ) as sync_client: + result: Final = completion(**args, client=sync_client) + text: Final = ( + "".join(chunk.choices[0].delta.content or "" for chunk in result) + if stream + else result.choices[0].message.content + ) + assert text == "OK" + sent: Final = requests.get_nowait() + assert str(sent.url) == f"{client_base}/chat/completions" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == { + "model": "session-test", + "messages": args["messages"], + **({"stream": True} if stream else {}), + } + assert requests.empty() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("source", ["OPENAI_BASE_URL", "OPENAI_API_BASE", "litellm.api_base"]) +@pytest.mark.parametrize("explicit_base", [None, "https://other.example/v1"]) +@pytest.mark.asyncio +async def test_completion_opencode_session_uses_effective_destination( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + asynchronous: bool, + source: str, + explicit_base: str | None, +) -> None: + default_base: Final = "https://opencode.ai/inference/go/openai/v1" + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", default_base if source == "litellm.api_base" else None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if source != "litellm.api_base": + monkeypatch.setenv(source, default_base) + endpoint: Final = respx_mock.post(f"{explicit_base or default_base}/chat/completions").respond( + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "session-test", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}], + } + ) + args: Final = { + "model": "openai/session-test", + "messages": [{"role": "user", "content": "hello"}], + "api_base": explicit_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + } + response: Final = await litellm.acompletion(**args) if asynchronous else litellm.completion(**args) + assert response.choices[0].message.content == "OK" + assert endpoint.call_count == 1 + sent: Final = endpoint.calls[0].request + assert sent.headers.get("x-opencode-session") == (None if explicit_base else "conversation-123") + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "session-test", "messages": args["messages"]} + + @pytest.fixture def openai_api_response(): mock_response_data = { @@ -3868,6 +4004,9 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): model="gpt-3.5-turbo-instruct", choices=[{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}], ), + TextCompletionResponse( + id="cmpl-1", object="text_completion", created=1, model="gpt-3.5-turbo-instruct", choices=[] + ), ] response = stream_chunk_builder_text_completion(chunks=chunks, messages=[{"role": "user", "content": "say hello"}]) @@ -8853,7 +8992,8 @@ async def test_text_completion_stream_include_usage_returns_usage_on_final_chunk "stream": True, "stream_options": {"include_usage": True}, } - assert "".join(chunk.choices[0].text or "" for chunk in chunks) == "reply" + assert "".join(chunk.choices[0].text or "" for chunk in chunks if chunk.choices) == "reply" + assert chunks[-1].choices == [] final_usage: Final = chunks[-1].usage assert (final_usage.prompt_tokens, final_usage.completion_tokens, final_usage.total_tokens) == (3, 2, 5) diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 701b79d0ba9..45222b2c9ab 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -27,6 +27,23 @@ def prices() -> dict: return json.loads(PRICES_PATH.read_text()) +@pytest.mark.parametrize("model", ("azure/gpt-audio-1.5-2026-02-23", "azure/gpt-audio-mini-2025-10-06")) +def test_azure_audio_snapshot_capabilities_match_declared_modalities(model: str, local_model_cost_map: None) -> None: + catalog: Final = TypeAdapter(dict[str, JsonValue]).validate_json(PRICES_PATH.read_bytes()) + backup: Final = TypeAdapter(dict[str, JsonValue]).validate_json(BACKUP_PRICES_PATH.read_bytes()) + entry: Final = catalog[model] + assert isinstance(entry, dict) + assert entry == backup[model] + input_modalities: Final = entry["supported_modalities"] + output_modalities: Final = entry["supported_output_modalities"] + assert isinstance(input_modalities, list) + assert isinstance(output_modalities, list) + expected: Final = ("audio" in input_modalities, "audio" in output_modalities) + info: Final = litellm.get_model_info(model) + assert (info["supports_audio_input"], info["supports_audio_output"]) == expected + assert (litellm.supports_audio_input(model), litellm.supports_audio_output(model)) == expected + + class _CopilotEndpointRow(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -455,30 +472,23 @@ def is_active_priced_mistral_chat_row(name: str, entry: Mapping[str, object]) -> ) -def cache_read_is_tenth_of_input(entry: Mapping[str, object]) -> bool: +def cache_read_is_below_input(entry: Mapping[str, object]) -> bool: cache_read: Final = entry.get("cache_read_input_token_cost") input_cost: Final = entry.get("input_cost_per_token") - return ( - isinstance(cache_read, float) - and isinstance(input_cost, (int, float)) - and 0 < cache_read < input_cost - and cache_read == pytest.approx(input_cost / 10) - ) + return isinstance(cache_read, float) and isinstance(input_cost, (int, float)) and 0 < cache_read < input_cost @pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) def test_active_mistral_chat_rows_price_cache_reads_below_input(path: Path): """A Mistral chat row without a cache-read rate bills cached prompt tokens at zero, so every - active priced row must carry one, and it must be cheaper than a fresh input token. Mistral - bills cached tokens at 10% of the input price for every model (docs.mistral.ai/studio/ - conversations/advanced/prompt-caching, read 2026-09-18), so the ratio is checked as well.""" + active priced row must carry one, and it must be cheaper than a fresh input token.""" rows: Mapping[str, object] = json.loads(path.read_text()) drifted: Final = [ f"{name}: cache_read={entry.get('cache_read_input_token_cost')} input={entry.get('input_cost_per_token')}" for name, entry in rows.items() if isinstance(entry, dict) and is_active_priced_mistral_chat_row(name, entry) - and not cache_read_is_tenth_of_input(entry) + and not cache_read_is_below_input(entry) ] assert drifted == [] @@ -791,3 +801,26 @@ def test_volcengine_doubao_seed_2_1_pro_260915_row_matches_published_metadata(pa sibling = rows["volcengine/doubao-seed-2-1-pro-260628"] assert isinstance(sibling, dict) assert all(entry[key] == sibling[key] for key in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost")) + + +def _bedrock_gpt_oss_rows(prices: Mapping[str, JsonValue]) -> tuple[str, ...]: + return tuple( + key + for key, row in prices.items() + if "gpt-oss" in key and isinstance(row, dict) and str(row.get("litellm_provider", "")).startswith("bedrock") + ) + + +@pytest.mark.parametrize("path", [PRICES_PATH, BACKUP_PRICES_PATH], ids=["main", "backup"]) +def test_bedrock_gpt_oss_rows_declare_vision_support_like_their_safeguard_siblings(path: Path) -> None: + prices: Final = json.loads(path.read_text()) + rows: Final = _bedrock_gpt_oss_rows(prices) + assert rows, "no Bedrock gpt-oss rows found" + reference: Final = prices["openai.gpt-oss-safeguard-120b"]["supports_vision"] + undeclared: Final = tuple(key for key in rows if prices[key].get("supports_vision") is not reference) + assert undeclared == (), f"modality routing only skips rows that declare supports_vision: {undeclared}" + + +def test_bedrock_gpt_oss_vision_is_explicitly_disabled_for_the_router_gate() -> None: + assert litellm.utils.is_vision_explicitly_disabled("bedrock/openai.gpt-oss-120b-1:0") is True + assert litellm.utils.is_vision_explicitly_disabled("bedrock_mantle/openai.gpt-oss-20b") is True diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 13ac55fa937..6aaa3379f47 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -5,16 +5,17 @@ import functools import gc import itertools import json +import linecache import logging import os import sys import threading import warnings -from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping from datetime import datetime, timedelta, timezone from pathlib import Path -from types import MappingProxyType, SimpleNamespace -from typing import TYPE_CHECKING, Final, Literal, cast +from types import CodeType, FrameType, MappingProxyType, SimpleNamespace +from typing import TYPE_CHECKING, Final, Literal, TypeVar, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -23,7 +24,7 @@ import pytest import respx from fastapi import HTTPException from opentelemetry import trace -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm import APIConnectionError, Router @@ -60,7 +61,10 @@ from litellm.router import ( _without_line_breaks, Span, ) +from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router_strategy import simple_shuffle +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.router_strategy.savings_baseline import Baseline, _models_in from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.router_utils.cooldown_handlers import ( _async_get_cooldown_deployments, @@ -93,7 +97,7 @@ from litellm.llms.base_llm.vector_store.transformation import( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) -from litellm.types.utils import CallTypes, CredentialItem +from litellm.types.utils import CallTypes, CredentialItem, Usage from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE @@ -108,6 +112,137 @@ if TYPE_CHECKING: from litellm.integrations.otel.logger import OpenTelemetryV2 +_ROUTER_DECISIONS_RESPONSE: Final[Mapping[str, object]] = { + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "damaged", "probability": 0.9}], + "usage": { + "input_tokens": 1, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 1, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2, + }, +} + + +_OPENAI_DECISIONS_PARAMS: Final[Mapping[str, object]] = { + "model": "openai/gpt-6-luna", + "api_base": "https://api.openai.com/v1", + "api_key": "k", +} +_BEDROCK_DECISIONS_PARAMS: Final[Mapping[str, object]] = { + "model": "bedrock/us.openai.gpt-5.6-terra", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", +} + + +def _decisions_leg(leg_id: str, params: Mapping[str, object], weight: int | None = None) -> dict[str, object]: + return { + "model_name": "decision-group", + "litellm_params": params if weight is None else {**params, "weight": weight}, + "model_info": {"id": leg_id}, + } + + +_OPENAI_DECISIONS_LEG: Final = _decisions_leg("openai-leg", _OPENAI_DECISIONS_PARAMS) +_BEDROCK_DECISIONS_LEG: Final = _decisions_leg("bedrock-leg", _BEDROCK_DECISIONS_PARAMS) + + +def _mixed_decisions_group_router(*legs: dict[str, object]) -> Router: + return Router(model_list=list(legs), routing_strategy="simple-shuffle") + + +@pytest.fixture +def _router_httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.asyncio +async def test_decisions_on_a_mixed_group_never_routes_to_a_leg_without_decisions_support( + respx_mock: respx.MockRouter, + _router_httpx_transport: None, +) -> None: + router: Final = _mixed_decisions_group_router(_OPENAI_DECISIONS_LEG, _BEDROCK_DECISIONS_LEG) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_ROUTER_DECISIONS_RESPONSE) + + for _ in range(20): + await router.adecisions( + model="decision-group", + input=[{"role": "user", "content": "The package arrived with a broken screen."}], + questions=[ + {"type": "predicate", "name": "damaged", "instructions": "Does the customer report a damaged item?"} + ], + ) + + assert route.call_count == 20 + + +def test_decisions_on_a_mixed_group_never_routes_to_a_leg_without_decisions_support_sync( + respx_mock: respx.MockRouter, +) -> None: + router: Final = _mixed_decisions_group_router(_OPENAI_DECISIONS_LEG, _BEDROCK_DECISIONS_LEG) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_ROUTER_DECISIONS_RESPONSE) + + for _ in range(20): + router.decisions( + model="decision-group", + input=[{"role": "user", "content": "The package arrived with a broken screen."}], + questions=[ + {"type": "predicate", "name": "damaged", "instructions": "Does the customer report a damaged item?"} + ], + ) + + assert route.call_count == 20 + + +@pytest.mark.asyncio +async def test_decisions_on_a_group_with_no_capable_leg_still_returns_the_provider_error() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "bedrock-only", + "litellm_params": { + "model": "bedrock/us.openai.gpt-5.6-terra", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + }, + "model_info": {"id": "bedrock-leg"}, + } + ] + ) + + with pytest.raises(litellm.BadRequestError, match="Unknown Decisions provider 'bedrock'"): + await router.adecisions( + model="bedrock-only", + input=[{"role": "user", "content": "The package arrived with a broken screen."}], + questions=[ + {"type": "predicate", "name": "damaged", "instructions": "Does the customer report a damaged item?"} + ], + ) + + +@pytest.mark.asyncio +async def test_chat_on_a_mixed_group_still_reaches_the_leg_without_decisions_support() -> None: + router: Final = _mixed_decisions_group_router( + _decisions_leg("openai-leg", _OPENAI_DECISIONS_PARAMS, weight=0), + _decisions_leg("bedrock-leg", _BEDROCK_DECISIONS_PARAMS, weight=1), + ) + + response: Final = await router.acompletion( + model="decision-group", + messages=[{"role": "user", "content": "Say hi"}], + mock_response="hi", + ) + + assert response._hidden_params["model_id"] == "bedrock-leg" + + def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): # initialize a real Router (env‑vars can be empty) router = litellm.Router( @@ -148,6 +283,15 @@ def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): assert kwargs["litellm_metadata"] == {"baz": 123} +def test_router_model_group_llm_provider_resolves_chatgpt_without_login(tmp_path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + params: Final = LiteLLM_Params(model="chatgpt/gpt-5.4") + + assert Router._model_group_llm_provider(params) == ("gpt-5.4", "chatgpt") + assert Router._model_group_llm_provider(LiteLLM_Params(model="no-such-provider-model-xyz")) == ("", "") + assert not any(tmp_path.iterdir()) + + def test_router_with_model_info_and_model_group(): """ Test edge case where user specifies model_group in model_info @@ -7889,6 +8033,29 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): router._count_pre_call_check_tokens(messages=None, input=None) +def test_count_pre_call_check_tokens_skips_the_history_a_compaction_item_replaces(): + router = litellm.Router(model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}]) + summary: Final = "The user planned a Tokyo trip." + compaction_item: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary}), + } + history: Final = [ + {"role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"role": "assistant", "content": "Day one, Asakusa. " * 200}, + ] + next_turn: Final = {"role": "user", "content": "And Osaka?"} + + compacted_tokens: Final = router._count_pre_call_check_tokens( + messages=None, input=[*history, compaction_item, next_turn] + ) + + assert compacted_tokens == router._count_pre_call_check_tokens( + messages=None, input=[{"role": "user", "content": summary}, next_turn] + ) + + def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): """ When neither messages nor input is provided (e.g. endpoints without prompt text), @@ -16726,6 +16893,7 @@ class _AnthropicMessagesFakeByteStream: self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} self.closed = False self.closed_as_superseded = False + self.close_contexts: tuple[bool, ...] = () def __aiter__(self): return self @@ -16736,6 +16904,7 @@ class _AnthropicMessagesFakeByteStream: return self._chunks.pop(0) async def aclose(self) -> None: + self.close_contexts = (*self.close_contexts, closing_superseded_attempt()) self.closed_as_superseded = self.closed_as_superseded or closing_superseded_attempt() self.closed = True @@ -17979,6 +18148,51 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] +class _AnthropicMessagesRetryStashRecorder: + def __init__(self) -> None: + self.observed: tuple[Usage | None, ...] = () + + def record(self, usage: Usage | None) -> None: + self.observed = (*self.observed, usage) + + +@pytest.mark.asyncio +async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage() -> None: + """The dropped attempt's usage stays available for terminal billing without reaching the retry's success log.""" + router: Final = _anthropic_messages_retry_router(num_retries=1) + logging_obj: Final = LiteLLMLogging( + model="glm", + messages=[{"role": "user", "content": "ping"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1), + litellm_call_id="retry-own-usage", + function_id="retry-own-usage", + ) + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) + stash_recorder: Final = _AnthropicMessagesRetryStashRecorder() + usage_adapter: Final = TypeAdapter(Usage | None) + + def retry_stream() -> _AnthropicMessagesFakeByteStream: + stash_recorder.record( + usage_adapter.validate_python(logging_obj.model_call_details.get("combined_usage_object")) + ) + return _anthropic_messages_retried_stream() + + provider: Final = _AnthropicMessagesScriptedProvider(_anthropic_messages_dropped_before_content, retry_stream) + + stream: Final = await _anthropic_messages_stream_through_router( + router, provider, litellm_logging_obj=logging_obj + ) + body: Final = tuple([chunk async for chunk in stream]) + + assert stash_recorder.observed == (None,) + assert body == (_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")) + assert "response_cost" not in logging_obj.model_call_details + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + + @pytest.mark.asyncio async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_opens(): """Both attempts log through the request's one logging object, so the superseded attempt's end-of-stream @@ -18001,6 +18215,7 @@ async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_o assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] assert superseded_closed_when_retry_opened == [True] + assert superseded.close_contexts[0] is True @pytest.mark.asyncio @@ -29670,6 +29885,73 @@ def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: py assert [d["model_info"]["id"] for d in filtered] == ["only"] +@pytest.mark.parametrize( + ("deployment_stream_options", "default_stream_options", "request_stream_options", "usage_chunk_choice_count"), + [ + ({"include_usage": True}, None, None, 1), + (None, {"include_usage": True}, None, 1), + ({"include_usage": True}, None, {"include_usage": True}, 0), + (None, None, {"include_usage": True}, 0), + ], + ids=["deployment-config", "router-default", "request-over-deployment", "request-only"], +) +@pytest.mark.asyncio +async def test_router_configured_stream_usage_keeps_a_choice_on_the_usage_chunk( + deployment_stream_options: dict[str, bool] | None, + default_stream_options: dict[str, bool] | None, + request_stream_options: dict[str, bool] | None, + usage_chunk_choice_count: int, +): + deployment_params: Final = {"model": "openai/gpt-5-mini", "api_key": "fake", "mock_response": "hi there"} + router: Final = Router( + model_list=[ + { + "model_name": "m", + "litellm_params": ( + deployment_params + if deployment_stream_options is None + else {**deployment_params, "stream_options": deployment_stream_options} + ), + } + ], + default_litellm_params=None if default_stream_options is None else {"stream_options": default_stream_options}, + ) + request_kwargs: Final = {} if request_stream_options is None else {"stream_options": request_stream_options} + + response: Final = await router.acompletion( + model="m", messages=[{"role": "user", "content": "hi"}], stream=True, **request_kwargs + ) + chunks: Final = [chunk async for chunk in response] + + assert getattr(chunks[-1], "usage", None) is not None + assert len(chunks[-1].choices) == usage_chunk_choice_count + assert all(len(chunk.choices) == 1 for chunk in chunks[:-1]) + + +@pytest.mark.parametrize( + ("stream_options", "marked_injected"), + [({"include_usage": True}, True), (None, False)], + ids=["include-usage", "no-stream-options"], +) +def test_update_kwargs_with_deployment_reads_stream_options_from_pydantic_litellm_params( + stream_options: dict[str, bool] | None, marked_injected: bool +): + from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD + from litellm.types.router import LiteLLM_Params + + kwargs: Final[dict[str, object]] = {"stream": True, "metadata": {}} + Router()._update_kwargs_with_deployment( + deployment={ + "model_name": "gpt-5-mini", + "litellm_params": LiteLLM_Params(model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options), + "model_info": {"id": "pydantic-params"}, + }, + kwargs=kwargs, + ) + + assert (kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True) is marked_injected + + @pytest.mark.asyncio async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens(): """After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone.""" @@ -29733,3 +30015,463 @@ async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_s wrapped.adopt_fallback_headers(fallback) assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + fallback_tokens + + +_SYNC_SECONDS: Final = 10 +_ROUTER_FILE: Final = Router.__init__.__code__.co_filename +_INDEX_LOCK_LINES: Final = frozenset( + number + for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1) + if "with self._deployment_index_lock" in line +) + +_Read = TypeVar("_Read") + + +def _concurrent_delete_deployment(index: int) -> dict[str, object]: + return { + "model_name": f"model-{index}", + "litellm_params": {"model": f"openai/gpt-{index}", "api_key": "k"}, + "model_info": { + "id": f"deployment-{index}", + "team_id": "team-1", + "team_public_model_name": "public-shared" if index >= 2 else f"public-{index}", + }, + } + + +def _concurrent_delete_router() -> Router: + return Router(model_list=[_concurrent_delete_deployment(index) for index in range(4)]) + + +def _read_during_held_write( + read: Callable[[], _Read], + write: Callable[[], object], + event: Literal["call", "c_call"] = "call", + write_after: CodeType | None = None, +) -> _Read: + """Run `read` on its own thread while `write` waits at its first `event` in the re-index, after the list shrank and + before the maps caught up. With `write_after`, the write starts when `read` returns from that function. The write + resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets the half-updated + list and a read that takes the lock waits for the write""" + reindex: Final = Router._update_deployment_indices_after_removal.__code__ + held: Final = threading.Event() + release: Final = threading.Event() + settled: Final = threading.Event() + released_in_time: Final[list[bool]] = [] + outcome: Final[list[_Read]] = [] + failure: Final[list[BaseException]] = [] + + def hold_in_reindex(frame: FrameType, observed: str, _arg: object) -> None: + if observed == event and frame.f_code is reindex and not held.is_set(): + held.set() + released_in_time.append(release.wait(timeout=_SYNC_SECONDS)) + + def run_write() -> None: + previous: Final = sys.getprofile() + sys.setprofile(hold_in_reindex) + try: + write() + finally: + sys.setprofile(previous) + + writer: Final = threading.Thread(target=run_write) + + def start_write() -> None: + writer.start() + assert held.wait(timeout=_SYNC_SECONDS), "the write never reached re-indexing" + + def watch_read(frame: FrameType, observed: str, _arg: object) -> Callable[..., object] | None: + if frame.f_code.co_filename != _ROUTER_FILE: + return None + if observed == "return" and frame.f_code is write_after and writer.ident is None: + start_write() + elif observed == "line" and held.is_set() and frame.f_lineno in _INDEX_LOCK_LINES: + settled.set() + return watch_read + + def run_read() -> None: + previous: Final = sys.gettrace() + sys.settrace(watch_read) + try: + outcome.append(read()) + except BaseException as error: + failure.append(error) + finally: + sys.settrace(previous) + settled.set() + + reader: Final = threading.Thread(target=run_read) + try: + if write_after is None: + start_write() + reader.start() + assert settled.wait(timeout=_SYNC_SECONDS), "the read neither returned nor reached the index lock" + finally: + release.set() + if reader.ident is not None: + reader.join() + if writer.ident is not None: + writer.join() + assert released_in_time == [True], "the write was not held for the whole read" + if failure: + raise failure[0] + return outcome[0] + + +def _deployment_ids(deployments: list[DeploymentTypedDict]) -> tuple[str, ...]: + return tuple(deployment["model_info"]["id"] for deployment in deployments) + + +@pytest.mark.parametrize( + ("read", "expected"), + [ + pytest.param( + lambda router: router.get_deployment("deployment-3").model_info.id, "deployment-3", id="get_deployment" + ), + pytest.param( + lambda router: router.get_model_info("deployment-3")["model_info"]["id"], + "deployment-3", + id="get_model_info", + ), + pytest.param( + lambda router: _deployment_ids(router.get_all_deployments("model-3")), + ("deployment-3",), + id="get_all_deployments", + ), + pytest.param( + lambda router: _deployment_ids(router.get_all_deployments("public-shared", team_id="team-1")), + ("deployment-2", "deployment-3"), + id="get_all_deployments-team", + ), + pytest.param( + lambda router: router.map_team_model("public-shared", "team-1"), "public-shared", id="map_team_model" + ), + pytest.param(lambda router: router.get_model_ids("model-3"), ["deployment-3"], id="get_model_ids-by-name"), + pytest.param( + lambda router: router.get_model_ids(), ["deployment-1", "deployment-2", "deployment-3"], id="get_model_ids" + ), + pytest.param( + lambda router: router.get_deployment_by_model_group_name("model-3").model_info.id, + "deployment-3", + id="get_deployment_by_model_group_name", + ), + pytest.param( + lambda router: router.get_model_listing_info("model-3").cost_map_keys, + ("openai/gpt-3",), + id="get_model_listing_info", + ), + pytest.param( + lambda router: router.team_model_has_alternatives("deployment-3"), True, id="team_model_has_alternatives" + ), + pytest.param( + lambda router: ( + router._get_model_group_deployment_usable_by_team( + model_group_name="model-3", team_id="team-1" + ).model_info.id + ), + "deployment-3", + id="model_group_deployment_usable_by_team", + ), + pytest.param( + lambda router: ( + router._get_team_public_name_deployment(model_id="public-shared", team_id="team-1").model_info.id + ), + "deployment-2", + id="team_public_name_deployment", + ), + pytest.param( + lambda router: _deployment_ids(router._team_deployments_across_teams("public-shared")), + ("deployment-2", "deployment-3"), + id="team_deployments_across_teams", + ), + pytest.param( + lambda router: _deployment_ids(list(router.deployments_named("model-3") or ())), + ("deployment-3",), + id="deployments_named", + ), + pytest.param( + lambda router: should_use_polling_for_request( + background_mode=True, + polling_via_cache_enabled=["openai"], + redis_cache=MagicMock(), + model="model-3", + llm_router=router, + ), + True, + id="response_polling_provider_lookup", + ), + pytest.param( + lambda router: _models_in(router, "model-3"), + (Baseline("openai/gpt-3", "deployment-3"),), + id="savings_baseline_candidates", + ), + ], +) +def test_lookup_during_a_concurrent_delete_returns_the_requested_deployment( + read: Callable[[Router], object], expected: object +) -> None: + router: Final = _concurrent_delete_router() + + observed: Final = _read_during_held_write(lambda: read(router), lambda: router.delete_deployment("deployment-0")) + + assert observed == expected + + +def test_lookup_during_an_upsert_returns_the_requested_deployment() -> None: + router: Final = _concurrent_delete_router() + rotated: Final = Deployment( + **{**_concurrent_delete_deployment(0), "litellm_params": {"model": "openai/gpt-0", "api_key": "rotated"}} + ) + + observed: Final = _read_during_held_write( + lambda: router.get_deployment("deployment-3").model_info.id, lambda: router.upsert_deployment(rotated) + ) + + assert observed == "deployment-3" + + +def test_add_during_a_concurrent_delete_indexes_the_added_deployment() -> None: + router: Final = _concurrent_delete_router() + + _read_during_held_write( + lambda: router.add_deployment(Deployment(**_concurrent_delete_deployment(4))), + lambda: router.delete_deployment("deployment-0"), + ) + + assert [router.get_deployment(f"deployment-{index}").model_info.id for index in range(1, 5)] == [ + f"deployment-{index}" for index in range(1, 5) + ] + + +def test_set_model_list_during_a_concurrent_delete_indexes_the_new_list() -> None: + router: Final = _concurrent_delete_router() + + _read_during_held_write( + lambda: router.set_model_list([_concurrent_delete_deployment(index) for index in range(10, 14)]), + lambda: router.delete_deployment("deployment-0"), + event="c_call", + ) + + assert router.get_model_ids() == [f"deployment-{index}" for index in range(10, 14)] + + +def test_team_model_has_alternatives_when_a_delete_starts_after_its_deployment_lookup() -> None: + router: Final = _concurrent_delete_router() + + has_alternatives: Final = _read_during_held_write( + lambda: router.team_model_has_alternatives("deployment-3"), + lambda: router.delete_deployment("deployment-0"), + write_after=Router.get_deployment.__code__, + ) + + assert has_alternatives is True + + +def test_init_adaptive_router_deployment_during_a_concurrent_delete_prices_the_listed_model() -> None: + priced: Final = { + **_concurrent_delete_deployment(3), + "litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06}, + } + router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) + adaptive: Final = Deployment( + model_name="adaptive", + litellm_params=LiteLLM_Params( + model="auto_router/adaptive_router", adaptive_router_config={"available_models": ["model-3"]} + ), + model_info=ModelInfo(id="adaptive"), + ) + + _read_during_held_write( + lambda: router.init_adaptive_router_deployment(deployment=adaptive), + lambda: router.delete_deployment("deployment-0"), + ) + + assert router.adaptive_routers["adaptive"][0].strategy.model_to_cost == {"model-3": 3e-06} + + +def test_complexity_router_adaptive_pool_during_a_concurrent_delete_prices_the_listed_model() -> None: + priced: Final = { + **_concurrent_delete_deployment(3), + "litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06}, + } + router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) + complexity: Final = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=router, + complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["model-3"]}}, + ) + + adaptive: Final = _read_during_held_write( + complexity._ensure_adaptive_router, lambda: router.delete_deployment("deployment-0") + ) + + assert adaptive is not None + assert adaptive.model_to_cost == {"model-3": 3e-06} + + +def test_delete_of_an_unknown_deployment_returns_none_and_keeps_the_list() -> None: + router: Final = _concurrent_delete_router() + + assert router.delete_deployment("deployment-9") is None + assert router.get_model_ids() == [f"deployment-{index}" for index in range(4)] + + +def test_get_model_ids_by_name_skips_team_deployments_only_when_asked() -> None: + shared: Final = {"model_name": "shared", "litellm_params": {"model": "openai/gpt-shared", "api_key": "k"}} + router: Final = Router( + model_list=[ + {**shared, "model_info": {"id": "global"}}, + {**shared, "model_info": {"id": "team", "team_id": "team-1", "team_public_model_name": "public"}}, + ] + ) + + assert router.get_model_ids("shared", exclude_team_models=True) == ["global"] + assert router.get_model_ids("shared") == ["global", "team"] + + +_IMAGE_PART = {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}} +_IMAGE_CHAT = [{"role": "user", "content": [{"type": "text", "text": "What is in this image?"}, _IMAGE_PART]}] +_IMAGE_RESPONSES_INPUT = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "What is in this image?"}, + {"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo="}, + ], + } +] +_VISION_LEG: Final[dict[str, object]] = {"litellm_params": {"model": "openai/gpt-5.6-terra"}, "model_info": {"id": "vision"}} +_DECLARED_TEXT_LEG: Final[dict[str, object]] = { + "litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, + "model_info": {"id": "declared-text", "supports_vision": False}, +} + + +def _mixed_router(*deployments: dict[str, object], routing_strategy: str = "latency-based-routing") -> litellm.Router: + return litellm.Router( + model_list=[{"model_name": "mixed", **deployment} for deployment in deployments], + routing_strategy=routing_strategy, + ) + + +def _candidate_ids( + router: litellm.Router, + messages: list[dict[str, object]] | None = None, + request_kwargs: dict[str, object] | None = None, +) -> list[str]: + candidates: Final = router._without_declared_text_only_deployments( + router.get_model_list(model_name="mixed"), "mixed", messages, request_kwargs + ) + return [d["model_info"]["id"] for d in candidates] + + +def test_router_picks_the_vision_leg_for_image_requests_without_enable_pre_call_checks(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + assert router.enable_pre_call_checks is False + for _ in range(5): + assert router.get_available_deployment(model="mixed", messages=_IMAGE_CHAT)["model_info"]["id"] == "vision" + + +@pytest.mark.asyncio +async def test_async_router_picks_the_vision_leg_for_responses_images_without_enable_pre_call_checks(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + for _ in range(5): + deployment: Final = await router.async_get_available_deployment( + model="mixed", request_kwargs={"input": _IMAGE_RESPONSES_INPUT} + ) + assert deployment["model_info"]["id"] == "vision" + + +def test_router_sends_an_image_to_the_text_only_leg_when_pre_call_checks_drop_the_vision_leg(): + router: Final = litellm.Router( + model_list=[ + {"model_name": "mixed", **_VISION_LEG, "model_info": {"id": "vision", "max_input_tokens": 1}}, + {"model_name": "mixed", **_DECLARED_TEXT_LEG}, + ], + routing_strategy="latency-based-routing", + enable_pre_call_checks=True, + ) + assert router.get_available_deployment(model="mixed", messages=_IMAGE_CHAT)["model_info"]["id"] == "declared-text" + + +def test_router_sends_an_image_to_the_text_only_leg_when_the_vision_leg_is_cooling_down(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + router.cooldown_cache.add_deployment_to_cooldown( + model_id="vision", original_exception=Exception("bad key"), exception_status=401, cooldown_time=60 + ) + assert router.get_available_deployment(model="mixed", messages=_IMAGE_CHAT)["model_info"]["id"] == "declared-text" + + +def test_without_declared_text_only_deployments_passes_a_chosen_deployment_through(): + assert litellm.Router._without_declared_text_only_deployments(_DECLARED_TEXT_LEG, "mixed", _IMAGE_CHAT, None) is _DECLARED_TEXT_LEG + + +def test_image_routing_drop_declared_text_only_legs_for_image_requests(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision"] + + +def test_image_routing_keep_declared_text_only_legs_for_text_requests(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + assert _candidate_ids(router, messages=[{"role": "user", "content": "Say hi"}]) == [ + "vision", + "declared-text", + ] + + +def test_image_routing_drop_declared_text_only_legs_for_responses_input_image(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + assert _candidate_ids(router, request_kwargs={"input": _IMAGE_RESPONSES_INPUT}) == ["vision"] + + +def test_image_routing_send_image_to_the_group_when_every_leg_is_declared_text_only(): + other_text_leg = {**_DECLARED_TEXT_LEG, "model_info": {"id": "declared-text-2", "supports_vision": False}} + router: Final = _mixed_router(_DECLARED_TEXT_LEG, other_text_leg) + assert _candidate_ids(router, messages=_IMAGE_CHAT) == [ + "declared-text", + "declared-text-2", + ] + + +def test_image_routing_keep_undeclared_custom_legs_for_image_requests(): + custom_leg = { + "litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct", "api_base": "http://localhost:8000"}, + "model_info": {"id": "custom"}, + } + router: Final = _mixed_router(_VISION_LEG, custom_leg) + assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision", "custom"] + + +def test_image_routing_resolve_vision_through_the_cost_map_and_base_model(): + catalog_text_leg = {"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {"id": "catalog-text"}} + azure_text_leg = { + "litellm_params": {"model": "azure/my-audio-deployment", "api_base": "https://x.openai.azure.com"}, + "model_info": {"id": "azure-text", "base_model": "azure/gpt-audio"}, + } + router: Final = _mixed_router(_VISION_LEG, catalog_text_leg, azure_text_leg) + assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision"] + + +@pytest.mark.parametrize( + ("deployment", "expected"), + [ + (_DECLARED_TEXT_LEG, True), + ({"litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"supports_vision": True}}, False), + ({"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {}}, True), + ({"litellm_params": {"model": "azure/my-deploy"}, "model_info": {"base_model": "azure/gpt-audio"}}, True), + ({"litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct"}, "model_info": {}}, False), + ({"litellm_params": {"model": "no-such-model-anywhere"}, "model_info": {}}, False), + ({"litellm_params": {}, "model_info": {}}, False), + ], + ids=["override-false", "override-true", "cost-map-false", "base-model-false", "undeclared", "unresolvable", "no-params"], +) +def test_deployment_declares_no_vision_reads_the_override_then_the_cost_map(deployment: dict[str, object], expected: bool): + assert litellm.Router._deployment_declares_no_vision(deployment, "mixed") is expected + + +def test_request_carries_image_sees_chat_parts_and_responses_input_items(): + assert litellm.Router._request_carries_image(_IMAGE_CHAT, {}) is True + assert litellm.Router._request_carries_image(None, {"input": _IMAGE_RESPONSES_INPUT}) is True + assert litellm.Router._request_carries_image([{"role": "user", "content": "Say hi"}], {}) is False + assert litellm.Router._request_carries_image(None, {"input": "Say hi"}) is False diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 7088818aa01..af8e7f53462 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1,8 +1,12 @@ +import asyncio import functools +import gc import json import os +import weakref from datetime import datetime, timezone -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -12,8 +16,12 @@ from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo +from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.types.litellm_params import RoutingStrategyName +from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( CallTypes, ModelResponse, @@ -1852,6 +1860,315 @@ def test_discard(model_list): assert len(litellm.service_callback) == 0 +_CALLBACK_LIST_NAMES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +def _callback_ids(list_name: str) -> list[int]: + return [id(callback) for callback in getattr(litellm, list_name)] + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("least-busy", "leastbusy_logger"), + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_releases_strategy_selector_across_router_lifecycles( + routing_strategy: RoutingStrategyName, selector_attr: str +): + def retire() -> weakref.ReferenceType[object]: + router: Final = Router(model_list=[], routing_strategy=routing_strategy) + selector: Final = getattr(router, selector_attr) + assert any(callback is selector for callback in litellm.callbacks) + router.discard() + return weakref.ref(selector) + + selector_refs: Final = tuple(retire() for _ in range(3)) + await asyncio.sleep(0) + gc.collect() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + assert all(selector_ref() is None for selector_ref in selector_refs) + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("least-busy", "leastbusy_logger"), + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_purges_selector_from_event_lists_after_serving_a_request( + routing_strategy: RoutingStrategyName, selector_attr: str +): + router: Final = Router( + model_list=[{"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}], + routing_strategy=routing_strategy, + ) + selector: Final = getattr(router, selector_attr) + await router.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + assert any(callback is selector for callback in litellm.success_callback) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def test_discard_keeps_callbacks_owned_by_other_live_routers(): + survivor: Final = Router(model_list=[], routing_strategy="latency-based-routing") + unrelated: Final = CustomLogger() + litellm.logging_callback_manager.add_litellm_callback(unrelated) + before_ids: Final = {list_name: _callback_ids(list_name) for list_name in _CALLBACK_LIST_NAMES} + retired: Final = Router(model_list=[], routing_strategy="least-busy") + assert any(callback is retired.leastbusy_logger for callback in litellm.callbacks) + + retired.discard() + + assert {list_name: _callback_ids(list_name) for list_name in before_ids} == before_ids + assert any(callback is survivor.lowestlatency_logger for callback in litellm.callbacks) + survivor.discard() + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_leaves_same_strategy_survivor_receiving_events( + routing_strategy: RoutingStrategyName, selector_attr: str +): + model_list: Final[list[DeploymentTypedDict]] = [ + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} + ] + retired: Final = Router(model_list=model_list, routing_strategy=routing_strategy) + survivor: Final = Router(model_list=model_list, routing_strategy=routing_strategy) + + retired.discard() + await survivor.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + + assert any(callback is getattr(survivor, selector_attr) for callback in litellm.success_callback) + survivor.discard() + + +def test_discard_drops_routing_group_selectors(): + router: Final = Router( + model_list=[ + {"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} + ], + routing_groups=[ + {"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "latency-based-routing"} + ], + ) + _, group_selector = router._get_routing_context("grouped-model") + assert any(callback is group_selector for callback in litellm.callbacks) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +_GROUPED_MODEL_LIST: Final[list[DeploymentTypedDict]] = [ + {"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} +] + + +class _LabelledHook(CustomLogger): + def __init__(self, label: str) -> None: + super().__init__() + self.label = label + + +def _visit_hooks_reloading_routing_groups_mid_loop( + router: Router, hooks: tuple[_LabelledHook, ...] +) -> Iterator[CustomLogger]: + for callback in litellm.callbacks: + if callback is hooks[0]: + router.update_settings( + routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}] + ) + if callback in hooks: + yield callback + + +def test_routing_group_reload_keeps_every_hook_of_an_in_flight_callback_loop(): + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}], + ) + hooks: Final = tuple(_LabelledHook(label) for label in ("guardrail_a", "guardrail_b", "guardrail_c")) + for hook in hooks: + litellm.logging_callback_manager.add_litellm_callback(hook) + + visited: Final = tuple(_visit_hooks_reloading_routing_groups_mid_loop(router, hooks)) + + assert visited == hooks + router.discard() + + +def _routers_seen_in_success_loop_discarding_first(first: Router, second: Router) -> Iterator[Router]: + for callback in litellm.success_callback: + if not isinstance(callback, MethodType): + continue + if callback.__self__ is first: + first.discard() + yield first + if callback.__self__ is second: + yield second + + +def test_discard_mid_success_loop_keeps_the_next_router_callback(): + first: Final = Router(model_list=_GROUPED_MODEL_LIST) + second: Final = Router(model_list=_GROUPED_MODEL_LIST) + + seen: Final = tuple(_routers_seen_in_success_loop_discarding_first(first, second)) + + assert seen == (first, second) + second.discard() + + +@pytest.mark.asyncio +async def test_discard_purges_optional_pre_call_checks_from_event_lists_after_serving_a_request(): + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + optional_pre_call_checks=["enforce_model_rate_limits", "prompt_caching", "deployment_affinity"], + ) + await router.acompletion(model="grouped-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def test_discard_registers_same_config_survivor_optional_pre_call_checks(): + retired: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"]) + survivor: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"]) + + retired.discard() + + assert survivor.optional_callbacks + assert all(_is_registered(check) for check in survivor.optional_callbacks) + survivor.discard() + + +def _is_registered(callback: object) -> bool: + return any(registered is callback for registered in litellm.callbacks) + + +@pytest.mark.asyncio +async def test_discard_releases_per_request_override_selector(): + def build_override_then_discard() -> weakref.ReferenceType[object]: + router: Final = Router(model_list=_GROUPED_MODEL_LIST) + selector: Final = router._get_override_strategy_selector("usage-based-routing-v2") + assert selector is not None + router.discard() + return weakref.ref(selector) + + selector_ref: Final = build_override_then_discard() + await asyncio.sleep(0) + gc.collect() + + assert selector_ref() is None + + + +def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list(): + retired: Final = Router(model_list=[], routing_strategy="latency-based-routing") + survivor: Final = Router(routing_strategy="latency-based-routing") + assert not _is_registered(survivor.lowestlatency_logger) + + retired.discard() + + assert _is_registered(survivor.lowestlatency_logger) + survivor.discard() + + +def test_discard_drops_its_slack_alerting_callbacks(): + router: Final = Router(model_list=_GROUPED_MODEL_LIST, alerting_config=AlertingConfig(webhook_url="test")) + assert any(isinstance(callback, SlackAlerting) for callback in litellm.callbacks) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def _adaptive_router_model_list(name: str) -> list[DeploymentTypedDict]: + return [ + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}, + { + "model_name": name, + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["gpt-5.5"]}, + }, + }, + ] + + +def _registered_adaptive_router_strategies() -> tuple[object, ...]: + return tuple( + callback.adaptive_router for callback in litellm.callbacks if isinstance(callback, AdaptiveRouterPostCallHook) + ) + + +def test_discard_drops_its_adaptive_router_hook(): + router: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + assert _registered_adaptive_router_strategies() + + router.discard() + + assert _registered_adaptive_router_strategies() == () + + +def test_discard_registers_the_adaptive_router_hook_of_a_survivor(): + survivor: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + retired: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + + retired.discard() + + survivor_strategies: Final = tuple(tagged.strategy for tagged in survivor.adaptive_routers["adaptive"]) + assert _registered_adaptive_router_strategies() == survivor_strategies + survivor.discard() + + +@pytest.mark.asyncio +async def test_discard_releases_its_router_budget_limiter(): + def build_budget_limiter_then_discard() -> weakref.ReferenceType[object]: + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + budget_limiter: Final = router.get_router_deployment_budget_limiter() + assert budget_limiter is not None + router.discard() + return weakref.ref(budget_limiter) + + budget_limiter_ref: Final = build_budget_limiter_then_discard() + await asyncio.sleep(0) + gc.collect() + + assert budget_limiter_ref() is None + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list) diff --git a/tests/unit/test_sync_together_ai_models.py b/tests/unit/test_sync_together_ai_models.py deleted file mode 100644 index b8a85bcfbdc..00000000000 --- a/tests/unit/test_sync_together_ai_models.py +++ /dev/null @@ -1,378 +0,0 @@ -import importlib.util -import json -from pathlib import Path -from types import MappingProxyType - -import pytest - -ROOT = Path(__file__).resolve().parents[2] -SCRIPT = ROOT / "scripts" / "sync_together_ai_models.py" -FIXTURES = Path(__file__).resolve().parent / "fixtures" / "together_ai_sync" - -_spec = importlib.util.spec_from_file_location("sync_together_ai_models", SCRIPT) -assert _spec is not None and _spec.loader is not None -sync = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(sync) - -RECORDED_CATALOG = sync.load_catalog(FIXTURES.joinpath("models_serverless.json").read_bytes()) -RECORDED_DOC = sync.parse_deprecations(FIXTURES.joinpath("deprecations.md").read_text()) - - -def _doc(removal_dates: dict[str, str], redirects: dict[str, str] | None = None) -> object: - return sync.DeprecationDoc( - removal_dates=MappingProxyType(removal_dates), - redirects=MappingProxyType(redirects or {}), - ) - - -def _chat_model(model_id: str, ctx: int = 4096, price: float = 1.0, cached: float | None = None) -> object: - return sync.CatalogModel( - id=model_id, - type="chat", - context_length=ctx, - pricing=sync.CatalogPricing(input=price, output=price, cached_input=cached), - ) - - -@pytest.mark.parametrize( - ("per_million", "expected"), - [ - (3, 3e-06), - (15, 1.5e-05), - (1.4, 1.4e-06), - (0.25999999999999995, 2.6e-07), - (0.060000000000000005, 6e-08), - (1.0399999999999998, 1.04e-06), - (0, 0.0), - ], -) -def test_per_token_normalizes_float_artifacts(per_million: float, expected: float) -> None: - assert sync.per_token(per_million) == expected - - -def test_parse_deprecations_recorded_fixture() -> None: - assert dict(RECORDED_DOC.redirects) == { - "mistralai/Mistral-7B-Instruct-v0.3": "mistralai/Ministral-3-14B-Instruct-2512", - "Kimi-K2": "Kimi-K2-0905", - "DeepSeek-V3": "DeepSeek-V3.1", - "DeepSeek-V3-0324": "DeepSeek-V3.1", - "DeepSeek-R1": "DeepSeek-R1-0528", - } - assert len(RECORDED_DOC.removal_dates) == 208 - assert RECORDED_DOC.removal_dates["google/gemma-3n-E4B-it"] == "2026-08-04" - - -def test_parse_deprecations_duplicate_rows_keep_most_recent_date() -> None: - assert RECORDED_DOC.removal_dates["Qwen/Qwen3-235B-A22B-Thinking-2507"] == "2026-04-16" - - -@pytest.mark.parametrize( - "markdown", - [ - "# Deprecations\n\nNothing here anymore.\n", - "\n## Active model redirects\n\n| A | B |\n| --- | --- |\n| `x` | `y` |\n\n## Something else\n", - "\n## Deprecation history\n\n### Inference\n\n| Date | Model | R |\n| --- | --- | --- |\n| 2026-01-01 | `m` | No |\n", - ], -) -def test_parse_deprecations_raises_when_a_table_parses_empty(markdown: str) -> None: - with pytest.raises(sync.SyncError): - sync.parse_deprecations(markdown) - - -def test_load_catalog_raises_on_shape_change() -> None: - with pytest.raises(sync.SyncError): - sync.load_catalog(b'[{"id": "x", "type": "chat"}]') - - -def test_load_catalog_raises_when_no_token_models_remain() -> None: - only_video = json.dumps([{"id": "v", "type": "video", "pricing": {"input": 0, "output": 0}}]).encode() - with pytest.raises(sync.SyncError): - sync.load_catalog(only_video) - - -def test_recorded_catalog_counts() -> None: - assert len(RECORDED_CATALOG) == 102 - assert sum(1 for model in RECORDED_CATALOG if model.type in sync.TYPE_TO_MODE) == 26 - assert sum(1 for model in RECORDED_CATALOG if model.pricing.cached_input) == 13 - - -def test_added_chat_model_matches_reviewed_registry_shape() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - assert len(outcome.added) == 26 - assert not outcome.deprecated - assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"] == { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 3e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048576, - "max_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "source": "https://docs.together.ai/docs/serverless-models", - "supports_function_calling": True, - "supports_parallel_function_calling": True, - "supports_prompt_caching": True, - "supports_reasoning": True, - "supports_response_schema": True, - "supports_tool_choice": True, - "supports_vision": True, - } - - -def test_added_embedding_model_has_no_output_token_cap() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - assert outcome.cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] == { - "input_cost_per_token": 2e-08, - "litellm_provider": "together_ai", - "max_input_tokens": 514, - "max_tokens": 514, - "mode": "embedding", - "output_cost_per_token": 2e-08, - "output_vector_size": 1024, - "source": "https://docs.together.ai/docs/serverless-models", - } - - -def test_moderation_type_maps_to_chat_mode() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - guard = outcome.cost_map["together_ai/meta-llama/Llama-Guard-4-12B"] - assert guard["mode"] == "chat" - assert "max_output_tokens" not in guard - - -def test_output_ceiling_comes_from_the_rule_never_from_context_length() -> None: - glm = next(model for model in RECORDED_CATALOG if model.id == "zai-org/GLM-5.2") - fresh = sync.compute_sync({}, [_chat_model("acme/unreviewed", ctx=1048576), glm], _doc({"x": "2026-01-01"})) - unreviewed = fresh.cost_map["together_ai/acme/unreviewed"] - assert "max_output_tokens" not in unreviewed - assert (unreviewed["max_input_tokens"], unreviewed["max_tokens"]) == (1048576, 1048576) - reviewed = fresh.cost_map["together_ai/zai-org/GLM-5.2"] - assert (reviewed["max_input_tokens"], reviewed["max_output_tokens"], reviewed["max_tokens"]) == ( - 1048575, - 128000, - 128000, - ) - inflated = { - "together_ai/zai-org/GLM-5.2": { - "input_cost_per_token": 1.4e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048575, - "max_output_tokens": 1048575, - "max_tokens": 1048575, - "mode": "chat", - "output_cost_per_token": 4.4e-06, - } - } - corrected = sync.compute_sync(inflated, [glm], _doc({"x": "2026-01-01"})) - assert corrected.cost_map["together_ai/zai-org/GLM-5.2"]["max_output_tokens"] == 128000 - assert any("max_output_tokens: 1048575 -> 128000" in line for line in corrected.updated) - - -def test_docs_removed_but_live_model_stays_live_with_warning() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - gemma = outcome.cost_map["together_ai/google/gemma-3n-E4B-it"] - assert "deprecation_date" not in gemma - assert any("gemma-3n-E4B-it" in warning and "2026-08-04" in warning for warning in outcome.warnings) - - -def test_price_change_updates_api_fields_and_keeps_curated_ones() -> None: - registry = { - "together_ai/acme/chat-1": { - "input_cost_per_token": 9e-07, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_token": 9e-07, - "supports_audio_input": True, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/chat-1", ctx=8192, price=2.0)], _doc({"x": "2026-01-01"})) - entry = outcome.cost_map["together_ai/acme/chat-1"] - assert entry["input_cost_per_token"] == 2e-06 - assert entry["max_input_tokens"] == 8192 - assert entry["max_output_tokens"] == 2048 - assert entry["supports_audio_input"] is True - assert len(outcome.updated) == 1 - assert "input_cost_per_token" in outcome.updated[0] - - -def test_cached_input_appearing_and_disappearing() -> None: - registry = { - "together_ai/acme/chat-1": { - "cache_read_input_token_cost": 1e-07, - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - "supports_prompt_caching": True, - }, - "together_ai/acme/chat-2": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - } - catalog = [_chat_model("acme/chat-1"), _chat_model("acme/chat-2", cached=0.25999999999999995)] - outcome = sync.compute_sync(registry, catalog, _doc({"x": "2026-01-01"})) - assert "cache_read_input_token_cost" not in outcome.cost_map["together_ai/acme/chat-1"] - assert "supports_prompt_caching" not in outcome.cost_map["together_ai/acme/chat-1"] - assert outcome.cost_map["together_ai/acme/chat-2"]["cache_read_input_token_cost"] == 2.6e-07 - assert outcome.cost_map["together_ai/acme/chat-2"]["supports_prompt_caching"] is True - - -def test_capability_rule_backfills_existing_entry() -> None: - registry = { - "together_ai/moonshotai/Kimi-K3": { - "input_cost_per_token": 3e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - } - } - kimi = next(model for model in RECORDED_CATALOG if model.id == "moonshotai/Kimi-K3") - outcome = sync.compute_sync(registry, [kimi], _doc({"x": "2026-01-01"})) - assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"]["supports_reasoning"] is True - assert any("supports_reasoning" in line for line in outcome.updated) - - -def test_disappeared_model_gets_docs_date_and_is_never_deleted() -> None: - registry = { - "together_ai/acme/gone": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"})) - assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-07-01" - assert outcome.deprecated == ("together_ai/acme/gone: deprecation_date",) - - -def test_disappeared_model_without_docs_date_warns_instead() -> None: - registry = { - "together_ai/acme/gone": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"other": "2026-07-01"})) - assert "deprecation_date" not in outcome.cost_map["together_ai/acme/gone"] - assert not outcome.deprecated - assert any("acme/gone" in warning and "human" in warning for warning in outcome.warnings) - - -def test_curated_deprecation_date_is_never_overwritten() -> None: - registry = { - "together_ai/acme/gone": { - "deprecation_date": "2026-06-15", - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"})) - assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-06-15" - assert any("2026-06-15" in warning and "2026-07-01" in warning for warning in outcome.warnings) - - -def test_redirect_chain_resolves_to_final_live_model() -> None: - doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b", "acme/b": "acme/c"}) - live = frozenset({"acme/c"}) - assert sync.resolve_successor("acme/a", doc, live) == "acme/c" - - -def test_redirect_dead_end_yields_no_successor() -> None: - doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b"}) - assert sync.resolve_successor("acme/a", doc, frozenset({"acme/other"})) is None - - -def test_redirect_short_names_resolve_by_unique_suffix() -> None: - doc = _doc({"moonshotai/Kimi-K2": "2026-01-01"}, redirects={"Kimi-K2": "Kimi-K2-0905"}) - live = frozenset({"moonshotai/Kimi-K2-0905"}) - assert sync.resolve_successor("moonshotai/Kimi-K2", doc, live) == "moonshotai/Kimi-K2-0905" - - -def test_successor_written_only_when_not_curated() -> None: - registry = { - "together_ai/acme/a": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - "together_ai/acme/b": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "metadata": {"successor": "together_ai/acme/curated"}, - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - } - doc = _doc({"acme/a": "2026-01-01", "acme/b": "2026-01-01"}, redirects={"acme/a": "acme/c", "acme/b": "acme/c"}) - outcome = sync.compute_sync(registry, [_chat_model("acme/c")], doc) - assert outcome.cost_map["together_ai/acme/a"]["metadata"] == {"successor": "together_ai/acme/c"} - assert outcome.cost_map["together_ai/acme/b"]["metadata"] == {"successor": "together_ai/acme/curated"} - assert any("acme/curated" in warning for warning in outcome.warnings) - - -def test_reappearance_clears_deprecation_date() -> None: - registry = { - "together_ai/acme/back": { - "deprecation_date": "2026-05-01", - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/back")], _doc({"x": "2026-01-01"})) - assert "deprecation_date" not in outcome.cost_map["together_ai/acme/back"] - assert outcome.reappeared == ("together_ai/acme/back",) - - -def test_new_chat_model_without_rule_is_flagged() -> None: - outcome = sync.compute_sync({}, [_chat_model("acme/unreviewed")], _doc({"x": "2026-01-01"})) - assert any("acme/unreviewed" in warning and "capability rule" in warning for warning in outcome.warnings) - - -def test_new_keys_land_at_the_end_of_the_provider_block() -> None: - registry = { - "aaa": {"mode": "chat"}, - "together_ai/acme/old": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - "zzz": {"mode": "chat"}, - } - outcome = sync.compute_sync(registry, [_chat_model("acme/old"), _chat_model("acme/new")], _doc({"x": "2026-01-01"})) - assert list(outcome.cost_map) == ["aaa", "together_ai/acme/old", "together_ai/acme/new", "zzz"] - - -def test_sync_is_idempotent_over_the_repo_cost_map() -> None: - cost_map = json.loads((ROOT / "model_prices_and_context_window.json").read_text()) - first = sync.compute_sync(cost_map, RECORDED_CATALOG, RECORDED_DOC) - second = sync.compute_sync(first.cost_map, RECORDED_CATALOG, RECORDED_DOC) - assert not second.has_changes - assert second.cost_map == first.cost_map - - -def test_pr_body_lists_every_section_and_the_skipped_types() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - body = sync.render_pr_body(outcome) - assert "### Added (26)" in body - assert "### Warnings needing a human call" in body - assert "image (29)" in body - assert "video (38)" in body diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 9b6700e9d39..b9baef4056a 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -21,7 +21,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest import respx -from jsonschema import validate +from jsonschema import Draft7Validator, validate import litellm from litellm.exceptions import MidStreamFallbackError @@ -90,7 +90,6 @@ from litellm.utils import ( get_llm_provider, get_optional_params_embeddings, ProviderConfigManager, - TextCompletionStreamWrapper, _check_provider_match, get_potential_model_names, _is_litellm_router_call, @@ -924,6 +923,457 @@ def validate_model_cost_values(model_data, exceptions=None): return len(violations) == 0, violations +_COST_MAP_ROW_SCHEMA: Final = { + "type": "object", + "additionalProperties": { + "type": "object", + "properties": { + "supports_computer_use": {"type": "boolean"}, + "cache_creation_input_audio_token_cost": {"type": "number"}, + "cache_creation_input_token_cost": {"type": "number"}, + "cache_creation_input_token_cost_above_1hr": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_flex": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_priority": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_ultrafast": {"type": "number"}, + "cache_creation_input_token_cost_above_32k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_100k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_100k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_128k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_200k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_256k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_flex": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_priority": {"type": "number"}, + "cache_creation_input_token_cost_above_200k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_batches": {"type": "number"}, + "cache_creation_input_token_cost_decisions": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, + "cache_read_input_token_cost_decisions": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, + "input_cost_per_token_decisions": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, + "output_cost_per_token_decisions": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, + "cache_creation_input_token_cost_flex": {"type": "number"}, + "cache_creation_input_token_cost_priority": {"type": "number"}, + "cache_creation_input_token_cost_ultrafast": {"type": "number"}, + "cache_read_input_token_cost": {"type": "number"}, + "explicit_cache_read_input_token_cost": {"type": "number"}, + "cache_read_input_token_cost_above_32k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_100k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_100k_tokens_batches": {"type": "number"}, + "cache_read_input_token_cost_above_128k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_200k_tokens_batches": {"type": "number"}, + "cache_read_input_token_cost_above_256k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_flex": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, + "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, + "cache_read_input_token_cost_batches": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": {"type": "number"}, + "cache_read_input_audio_token_cost": {"type": "number"}, + "cache_read_input_audio_token_cost_batches": {"type": "number"}, + "cache_read_input_audio_token_cost_flex": {"type": "number"}, + "cache_read_input_audio_token_cost_priority": {"type": "number"}, + "cache_read_input_image_token_cost": {"type": "number"}, + "audio_transcription_config": {"type": "string"}, + "deprecation_date": {"type": "string"}, + "input_cost_per_audio_per_second": {"type": "number"}, + "input_cost_per_audio_per_second_above_128k_tokens": {"type": "number"}, + "google_maps_grounding_cost_per_query": {"type": "number"}, + "input_cost_per_audio_token": {"type": "number"}, + "input_cost_per_audio_token_flex": {"type": "number"}, + "input_cost_per_image_token": {"type": "number"}, + "input_cost_per_character": {"type": "number"}, + "input_cost_per_character_above_128k_tokens": {"type": "number"}, + "input_cost_per_image": {"type": "number"}, + "input_cost_per_image_above_128k_tokens": {"type": "number"}, + "input_cost_per_video_token": {"type": "number"}, + "input_cost_per_token_above_32k_tokens": {"type": "number"}, + "input_cost_per_token_above_100k_tokens": {"type": "number"}, + "input_cost_per_token_above_100k_tokens_batches": {"type": "number"}, + "input_cost_per_token_above_200k_tokens": {"type": "number"}, + "input_cost_per_token_above_200k_tokens_batches": {"type": "number"}, + "input_cost_per_token_above_256k_tokens": {"type": "number"}, + "input_cost_per_token_above_272k_tokens": {"type": "number"}, + "input_cost_per_token_above_512k_tokens": {"type": "number"}, + "cache_read_input_token_cost_flex": {"type": "number"}, + "cache_read_input_token_cost_priority": {"type": "number"}, + "cache_read_input_token_cost_balanced": {"type": "number"}, + "cache_read_input_token_cost_ultrafast": {"type": "number"}, + "cache_read_input_token_cost_above_200k_tokens_priority": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_priority": {"type": "number"}, + "input_cost_per_token_flex": {"type": "number"}, + "input_cost_per_token_priority": {"type": "number"}, + "input_cost_per_token_balanced": {"type": "number"}, + "input_cost_per_token_ultrafast": {"type": "number"}, + "input_cost_per_token_above_200k_tokens_priority": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_priority": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_batches": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_flex": {"type": "number"}, + "input_cost_per_audio_token_priority": {"type": "number"}, + "output_cost_per_token_flex": {"type": "number"}, + "output_cost_per_token_priority": {"type": "number"}, + "output_cost_per_token_balanced": {"type": "number"}, + "output_cost_per_token_ultrafast": {"type": "number"}, + "output_cost_per_token_above_200k_tokens_priority": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_priority": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_batches": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_flex": {"type": "number"}, + "regional_endpoint_uplift_multiplier": {"type": "number"}, + "regional_processing_uplift_multiplier_eu": {"type": "number"}, + "regional_processing_uplift_multiplier_us": {"type": "number"}, + "input_cost_per_pixel": {"type": "number"}, + "input_cost_per_query": {"type": "number"}, + "input_cost_per_request": {"type": "number"}, + "cost_per_second": {"type": "number"}, + "input_cost_per_second": {"type": "number"}, + "input_cost_per_text_message": {"type": "number"}, + "input_cost_per_token": {"type": "number"}, + "input_cost_per_token_above_128k_tokens": {"type": "number"}, + "input_cost_per_audio_token_batches": {"type": "number"}, + "input_cost_per_image_token_batches": {"type": "number"}, + "input_cost_per_token_batches": {"type": "number"}, + "input_cost_per_video_token_batches": {"type": "number"}, + "input_cost_per_token_cache_hit": {"type": "number"}, + "input_cost_per_video_per_second": {"type": "number"}, + "input_cost_per_video_per_second_above_8s_interval": {"type": "number"}, + "input_cost_per_video_per_second_above_15s_interval": {"type": "number"}, + "input_cost_per_video_per_second_above_128k_tokens": {"type": "number"}, + "input_dbu_cost_per_token": {"type": "number"}, + "annotation_cost_per_page": {"type": "number"}, + "annotation_cost_per_page_batches": {"type": "number"}, + "ocr_cost_per_page": {"type": "number"}, + "ocr_cost_per_page_batches": {"type": "number"}, + "ocr_cost_per_credit": {"type": "number"}, + "code_interpreter_cost_per_session": {"type": "number"}, + "inference_geo": {"type": "string"}, + "litellm_provider": {"type": "string"}, + "max_input_tokens": {"type": "number"}, + "max_output_tokens": {"type": "number"}, + "max_tokens": {"type": "number"}, + "metadata": {"type": "object"}, + "provider_specific_entry": {"type": "object"}, + "mode": { + "type": "string", + "enum": [ + "audio_speech", + "audio_transcription", + "chat", + "completion", + "container", + "decisions", + "image_edit", + "embedding", + "evaluation", + "extraction", + "guardrail", + "image_generation", + "video_generation", + "moderation", + "rerank", + "realtime", + "responses", + "ocr", + "search", + "vector_store", + ], + }, + "output_cost_per_audio_token": {"type": "number"}, + "output_cost_per_character": {"type": "number"}, + "output_cost_per_character_above_128k_tokens": {"type": "number"}, + "output_cost_per_image": {"type": "number"}, + "output_cost_per_image_512": {"type": "number"}, + "output_cost_per_image_1024": {"type": "number"}, + "output_cost_per_image_1536": {"type": "number"}, + "output_cost_per_image_0.5K": {"type": "number"}, + "output_cost_per_image_1K": {"type": "number"}, + "output_cost_per_image_2K": {"type": "number"}, + "output_cost_per_image_4K": {"type": "number"}, + "output_cost_per_image_token": {"type": "number"}, + "output_cost_per_image_token_batches": {"type": "number"}, + "output_cost_per_video_token": {"type": "number"}, + "output_cost_per_pixel": {"type": "number"}, + "output_cost_per_second": {"type": "number"}, + "output_cost_per_second_480p": {"type": "number"}, + "output_cost_per_second_720p": {"type": "number"}, + "output_cost_per_second_768p": {"type": "number"}, + "output_cost_per_second_2k": {"type": "number"}, + "output_cost_per_second_1080p": {"type": "number"}, + "output_cost_per_second_4k": {"type": "number"}, + "output_cost_per_token": {"type": "number"}, + "output_cost_per_token_above_32k_tokens": {"type": "number"}, + "output_cost_per_token_above_100k_tokens": {"type": "number"}, + "output_cost_per_token_above_100k_tokens_batches": {"type": "number"}, + "output_cost_per_token_above_128k_tokens": {"type": "number"}, + "output_cost_per_token_above_200k_tokens": {"type": "number"}, + "output_cost_per_token_above_200k_tokens_batches": {"type": "number"}, + "output_cost_per_token_above_256k_tokens": {"type": "number"}, + "output_cost_per_token_above_272k_tokens": {"type": "number"}, + "output_cost_per_token_above_512k_tokens": {"type": "number"}, + "output_cost_per_token_batches": {"type": "number"}, + "output_cost_per_reasoning_token": {"type": "number"}, + "output_cost_per_video_per_second": {"type": "number"}, + "output_db_cost_per_token": {"type": "number"}, + "output_dbu_cost_per_token": {"type": "number"}, + "output_vector_size": {"type": "number"}, + "rpd": {"type": "number"}, + "rpm": {"type": "number"}, + "source": {"type": "string"}, + "comment": {"type": "string"}, + "supports_assistant_prefill": {"type": "boolean"}, + "supports_anthropic_compaction": {"type": "boolean"}, + "supports_audio_input": {"type": "boolean"}, + "supports_audio_output": {"type": "boolean"}, + "gemini_native_audio": {"type": "boolean"}, + "gemini_audio_only_live": {"type": "boolean"}, + "supports_embedding_image_input": {"type": "boolean"}, + "supports_forced_tool_use": {"type": "boolean"}, + "supports_function_calling": {"type": "boolean"}, + "supports_image_input": {"type": "boolean"}, + "supports_nova_canvas_image_edit": {"type": "boolean"}, + "supports_parallel_function_calling": {"type": "boolean"}, + "supports_parallel_tool_use_config": {"type": "boolean"}, + "supports_pdf_input": {"type": "boolean"}, + "prompt_cache_min_tokens": {"type": "number"}, + "supports_prompt_cache_breakpoint": {"type": "boolean"}, + "supports_thinking_cache_preservation": {"type": "boolean"}, + "supports_prompt_caching": {"type": "boolean"}, + "supports_response_schema": {"type": "boolean"}, + "supports_system_messages": {"type": "boolean"}, + "supports_tool_choice": {"type": "boolean"}, + "supports_tool_search": {"type": "boolean"}, + "supports_video_input": {"type": "boolean"}, + "supports_box_questions": {"type": "boolean"}, + "supports_vision": {"type": "boolean"}, + "supports_web_search": {"type": "boolean"}, + "token_count_framing": {"type": "integer"}, + "token_count_model": { + "type": "object", + "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, + "required": ["model"], + }, + "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]}, + "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"}, + "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"}, + "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, + "supports_bedrock_converse_tool_result_images": {"type": "boolean"}, + "supports_url_context": {"type": "boolean"}, + "supports_multimodal": {"type": "boolean"}, + "uses_embed_content": {"type": "boolean"}, + "supports_reasoning": {"type": "boolean"}, + "supports_minimal_reasoning_effort": {"type": "boolean"}, + "supports_low_reasoning_effort": {"type": "boolean"}, + "supports_none_reasoning_effort": {"type": "boolean"}, + "supports_xhigh_reasoning_effort": {"type": "boolean"}, + "supports_max_reasoning_effort": {"type": "boolean"}, + "reasoning_effort_levels": { + "type": "array", + "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, + }, + "default_reasoning_effort": { + "type": "string", + "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], + }, + "supports_adaptive_thinking": {"type": "boolean"}, + "supports_anthropic_thinking_payload": {"type": "boolean"}, + "supports_legacy_thinking": {"type": "boolean"}, + "thinking_always_on": {"type": "boolean"}, + "supports_mid_conversation_system": {"type": "boolean"}, + "supports_sampling_params": {"type": "boolean"}, + "supports_output_config": {"type": "boolean"}, + "supports_speed": {"type": "boolean"}, + "supports_fast_mode": {"type": "boolean"}, + "supported_audio_formats": { + "type": "array", + "items": { + "type": "string", + "enum": ["mp3", "wav"], + }, + }, + "vertex_ai_audio_api": { + "type": "string", + "enum": ["lyria_predict", "lyria_interactions"], + }, + "bedrock_output_config_effort_ceiling": { + "type": "string", + "enum": ["low", "medium", "high", "max", "xhigh"], + }, + "bedrock_converse_supports_strict_tools": {"type": "boolean"}, + "supports_regex_lookaround": {"type": "boolean"}, + "tpm": {"type": "number"}, + "supported_endpoints": { + "type": "array", + "items": { + "type": "string", + "enum": [ + "/v1/responses", + "/v1/embeddings", + "/v1/chat/completions", + "/v1/completions", + "/v1/messages", + "/v1/images/generations", + "/v1/realtime", + "/v1/realtime/transcription_sessions", + "/v1/images/variations", + "/v1/images/edits", + "/v1/batch", + "/v1beta/interactions", + "/v1beta/models", + "/v1/audio/transcriptions", + "/v1/audio/speech", + "/v1/ocr", + "/v1/videos", + "/vertex_ai/live", + "/v1/listen", + "/v1/systemone", + "/v1/decisions", + "/v1beta/interactions", + ], + }, + }, + "supported_regions": { + "type": "array", + "items": { + "type": "string", + }, + }, + "guardrail_cost_per_unit": { + "type": "object", + "additionalProperties": {"type": "number"}, + }, + "search_context_cost_per_query": { + "type": "object", + "properties": { + "search_context_size_low": {"type": "number"}, + "search_context_size_medium": {"type": "number"}, + "search_context_size_high": {"type": "number"}, + }, + "additionalProperties": False, + }, + "web_search_billing_unit": { + "type": "string", + "enum": ["per_prompt", "per_query"], + }, + "search_free_requests": { + "type": "object", + "properties": { + "count": {"type": "integer", "minimum": 1}, + "period": {"type": "string", "enum": ["day", "month"]}, + "shared_group": {"type": "string"}, + }, + "required": ["count", "period", "shared_group"], + "additionalProperties": False, + }, + "citation_cost_per_token": {"type": "number"}, + "supported_modalities": { + "type": "array", + "items": { + "type": "string", + "enum": ["text", "audio", "image", "video"], + }, + }, + "supported_output_modalities": { + "type": "array", + "items": { + "type": "string", + "enum": ["text", "image", "audio", "code", "video"], + }, + }, + "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", + "properties": { + "hours_utc": { + "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], + }, + "windows": { + "type": "array", + "items": { + "type": "object", + "properties": { + "hours_utc": { + "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], + }, + "weekdays": { + "type": "array", + "items": {"oneOf": [{"type": "integer"}, {"type": "string"}]}, + }, + }, + "required": ["hours_utc"], + "additionalProperties": False, + }, + }, + "weekday_timezone": {"type": "string"}, + "input_cost_per_token": {"type": "number"}, + "output_cost_per_token": {"type": "number"}, + "output_cost_per_reasoning_token": {"type": "number"}, + "cache_read_input_token_cost": {"type": "number"}, + "cache_creation_input_token_cost": {"type": "number"}, + }, + "additionalProperties": False, + }, + "tiered_pricing": { + "type": "array", + "items": { + "type": "object", + "properties": { + "range": { + "type": "array", + "items": {"type": "number"}, + "minItems": 2, + "maxItems": 2, + }, + "input_cost_per_token": {"type": "number"}, + "output_cost_per_token": {"type": "number"}, + "cache_read_input_token_cost": {"type": "number"}, + "cache_creation_input_token_cost": {"type": "number"}, + "output_cost_per_reasoning_token": {"type": "number"}, + "max_results_range": { + "type": "array", + "items": {"type": "number"}, + "minItems": 2, + "maxItems": 2, + }, + "input_cost_per_query": {"type": "number"}, + }, + "additionalProperties": False, + }, + }, + }, + "additionalProperties": False, + }, +} + + def test_aaamodel_prices_and_context_window_json_is_valid(): """ Validates the `model_prices_and_context_window.json` file. @@ -931,430 +1381,6 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): If this test fails after you update the json, you need to update the schema or correct the change you made. """ - INTENDED_SCHEMA = { - "type": "object", - "additionalProperties": { - "type": "object", - "properties": { - "supports_computer_use": {"type": "boolean"}, - "cache_creation_input_audio_token_cost": {"type": "number"}, - "cache_creation_input_token_cost": {"type": "number"}, - "cache_creation_input_token_cost_above_1hr": {"type": "number"}, - "cache_creation_input_token_cost_above_32k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_100k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_100k_tokens_batches": {"type": "number"}, - "cache_creation_input_token_cost_above_128k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_200k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_256k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_flex": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_priority": {"type": "number"}, - "cache_creation_input_token_cost_above_200k_tokens_batches": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_batches": {"type": "number"}, - "cache_creation_input_token_cost_batches": {"type": "number"}, - "cache_creation_input_token_cost_decisions": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, - "cache_read_input_token_cost_decisions": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, - "input_cost_per_token_decisions": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, - "output_cost_per_token_decisions": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, - "cache_creation_input_token_cost_flex": {"type": "number"}, - "cache_creation_input_token_cost_priority": {"type": "number"}, - "cache_creation_input_token_cost_ultrafast": {"type": "number"}, - "cache_read_input_token_cost": {"type": "number"}, - "explicit_cache_read_input_token_cost": {"type": "number"}, - "cache_read_input_token_cost_above_32k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_100k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_100k_tokens_batches": {"type": "number"}, - "cache_read_input_token_cost_above_128k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_200k_tokens_batches": {"type": "number"}, - "cache_read_input_token_cost_above_256k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_flex": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, - "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, - "cache_read_input_token_cost_batches": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_batches": {"type": "number"}, - "cache_creation_input_token_cost_above_1hr_above_100k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_1hr_above_200k_tokens": {"type": "number"}, - "cache_read_input_audio_token_cost": {"type": "number"}, - "cache_read_input_image_token_cost": {"type": "number"}, - "audio_transcription_config": {"type": "string"}, - "deprecation_date": {"type": "string"}, - "input_cost_per_audio_per_second": {"type": "number"}, - "input_cost_per_audio_per_second_above_128k_tokens": {"type": "number"}, - "google_maps_grounding_cost_per_query": {"type": "number"}, - "input_cost_per_audio_token": {"type": "number"}, - "input_cost_per_image_token": {"type": "number"}, - "input_cost_per_character": {"type": "number"}, - "input_cost_per_character_above_128k_tokens": {"type": "number"}, - "input_cost_per_image": {"type": "number"}, - "input_cost_per_image_above_128k_tokens": {"type": "number"}, - "input_cost_per_video_token": {"type": "number"}, - "input_cost_per_token_above_32k_tokens": {"type": "number"}, - "input_cost_per_token_above_100k_tokens": {"type": "number"}, - "input_cost_per_token_above_100k_tokens_batches": {"type": "number"}, - "input_cost_per_token_above_200k_tokens": {"type": "number"}, - "input_cost_per_token_above_200k_tokens_batches": {"type": "number"}, - "input_cost_per_token_above_256k_tokens": {"type": "number"}, - "input_cost_per_token_above_272k_tokens": {"type": "number"}, - "input_cost_per_token_above_512k_tokens": {"type": "number"}, - "cache_read_input_token_cost_flex": {"type": "number"}, - "cache_read_input_token_cost_priority": {"type": "number"}, - "cache_read_input_token_cost_balanced": {"type": "number"}, - "cache_read_input_token_cost_ultrafast": {"type": "number"}, - "cache_read_input_token_cost_above_200k_tokens_priority": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_priority": {"type": "number"}, - "input_cost_per_token_flex": {"type": "number"}, - "input_cost_per_token_priority": {"type": "number"}, - "input_cost_per_token_balanced": {"type": "number"}, - "input_cost_per_token_ultrafast": {"type": "number"}, - "input_cost_per_token_above_200k_tokens_priority": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_priority": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_batches": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_flex": {"type": "number"}, - "input_cost_per_audio_token_priority": {"type": "number"}, - "output_cost_per_token_flex": {"type": "number"}, - "output_cost_per_token_priority": {"type": "number"}, - "output_cost_per_token_balanced": {"type": "number"}, - "output_cost_per_token_ultrafast": {"type": "number"}, - "output_cost_per_token_above_200k_tokens_priority": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_priority": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_batches": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_flex": {"type": "number"}, - "regional_endpoint_uplift_multiplier": {"type": "number"}, - "regional_processing_uplift_multiplier_eu": {"type": "number"}, - "regional_processing_uplift_multiplier_us": {"type": "number"}, - "input_cost_per_pixel": {"type": "number"}, - "input_cost_per_query": {"type": "number"}, - "input_cost_per_request": {"type": "number"}, - "cost_per_second": {"type": "number"}, - "input_cost_per_second": {"type": "number"}, - "input_cost_per_text_message": {"type": "number"}, - "input_cost_per_token": {"type": "number"}, - "input_cost_per_token_above_128k_tokens": {"type": "number"}, - "input_cost_per_audio_token_batches": {"type": "number"}, - "input_cost_per_image_token_batches": {"type": "number"}, - "input_cost_per_token_batches": {"type": "number"}, - "input_cost_per_video_token_batches": {"type": "number"}, - "input_cost_per_token_cache_hit": {"type": "number"}, - "input_cost_per_video_per_second": {"type": "number"}, - "input_cost_per_video_per_second_above_8s_interval": {"type": "number"}, - "input_cost_per_video_per_second_above_15s_interval": {"type": "number"}, - "input_cost_per_video_per_second_above_128k_tokens": {"type": "number"}, - "input_dbu_cost_per_token": {"type": "number"}, - "annotation_cost_per_page": {"type": "number"}, - "annotation_cost_per_page_batches": {"type": "number"}, - "ocr_cost_per_page": {"type": "number"}, - "ocr_cost_per_page_batches": {"type": "number"}, - "ocr_cost_per_credit": {"type": "number"}, - "code_interpreter_cost_per_session": {"type": "number"}, - "inference_geo": {"type": "string"}, - "litellm_provider": {"type": "string"}, - "max_input_tokens": {"type": "number"}, - "max_output_tokens": {"type": "number"}, - "max_tokens": {"type": "number"}, - "metadata": {"type": "object"}, - "provider_specific_entry": {"type": "object"}, - "mode": { - "type": "string", - "enum": [ - "audio_speech", - "audio_transcription", - "chat", - "completion", - "container", - "image_edit", - "embedding", - "evaluation", - "extraction", - "guardrail", - "image_generation", - "video_generation", - "moderation", - "rerank", - "realtime", - "responses", - "ocr", - "search", - "vector_store", - ], - }, - "output_cost_per_audio_token": {"type": "number"}, - "output_cost_per_character": {"type": "number"}, - "output_cost_per_character_above_128k_tokens": {"type": "number"}, - "output_cost_per_image": {"type": "number"}, - "output_cost_per_image_512": {"type": "number"}, - "output_cost_per_image_1024": {"type": "number"}, - "output_cost_per_image_1536": {"type": "number"}, - "output_cost_per_image_0.5K": {"type": "number"}, - "output_cost_per_image_1K": {"type": "number"}, - "output_cost_per_image_2K": {"type": "number"}, - "output_cost_per_image_4K": {"type": "number"}, - "output_cost_per_image_token": {"type": "number"}, - "output_cost_per_image_token_batches": {"type": "number"}, - "output_cost_per_video_token": {"type": "number"}, - "output_cost_per_pixel": {"type": "number"}, - "output_cost_per_second": {"type": "number"}, - "output_cost_per_second_480p": {"type": "number"}, - "output_cost_per_second_720p": {"type": "number"}, - "output_cost_per_second_768p": {"type": "number"}, - "output_cost_per_second_2k": {"type": "number"}, - "output_cost_per_second_1080p": {"type": "number"}, - "output_cost_per_second_4k": {"type": "number"}, - "output_cost_per_token": {"type": "number"}, - "output_cost_per_token_above_32k_tokens": {"type": "number"}, - "output_cost_per_token_above_100k_tokens": {"type": "number"}, - "output_cost_per_token_above_100k_tokens_batches": {"type": "number"}, - "output_cost_per_token_above_128k_tokens": {"type": "number"}, - "output_cost_per_token_above_200k_tokens": {"type": "number"}, - "output_cost_per_token_above_200k_tokens_batches": {"type": "number"}, - "output_cost_per_token_above_256k_tokens": {"type": "number"}, - "output_cost_per_token_above_272k_tokens": {"type": "number"}, - "output_cost_per_token_above_512k_tokens": {"type": "number"}, - "output_cost_per_token_batches": {"type": "number"}, - "output_cost_per_reasoning_token": {"type": "number"}, - "output_cost_per_video_per_second": {"type": "number"}, - "output_db_cost_per_token": {"type": "number"}, - "output_dbu_cost_per_token": {"type": "number"}, - "output_vector_size": {"type": "number"}, - "rpd": {"type": "number"}, - "rpm": {"type": "number"}, - "source": {"type": "string"}, - "comment": {"type": "string"}, - "supports_assistant_prefill": {"type": "boolean"}, - "supports_anthropic_compaction": {"type": "boolean"}, - "supports_audio_input": {"type": "boolean"}, - "supports_audio_output": {"type": "boolean"}, - "gemini_native_audio": {"type": "boolean"}, - "gemini_audio_only_live": {"type": "boolean"}, - "supports_embedding_image_input": {"type": "boolean"}, - "supports_forced_tool_use": {"type": "boolean"}, - "supports_function_calling": {"type": "boolean"}, - "supports_image_input": {"type": "boolean"}, - "supports_nova_canvas_image_edit": {"type": "boolean"}, - "supports_parallel_function_calling": {"type": "boolean"}, - "supports_parallel_tool_use_config": {"type": "boolean"}, - "supports_pdf_input": {"type": "boolean"}, - "prompt_cache_min_tokens": {"type": "number"}, - "supports_prompt_cache_breakpoint": {"type": "boolean"}, - "supports_thinking_cache_preservation": {"type": "boolean"}, - "supports_prompt_caching": {"type": "boolean"}, - "supports_response_schema": {"type": "boolean"}, - "supports_system_messages": {"type": "boolean"}, - "supports_tool_choice": {"type": "boolean"}, - "supports_tool_search": {"type": "boolean"}, - "supports_video_input": {"type": "boolean"}, - "supports_box_questions": {"type": "boolean"}, - "supports_vision": {"type": "boolean"}, - "supports_web_search": {"type": "boolean"}, - "token_count_framing": {"type": "integer"}, - "token_count_model": { - "type": "object", - "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, - "required": ["model"], - }, - "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]}, - "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"}, - "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"}, - "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, - "supports_bedrock_converse_tool_result_images": {"type": "boolean"}, - "supports_url_context": {"type": "boolean"}, - "supports_multimodal": {"type": "boolean"}, - "uses_embed_content": {"type": "boolean"}, - "supports_reasoning": {"type": "boolean"}, - "supports_minimal_reasoning_effort": {"type": "boolean"}, - "supports_low_reasoning_effort": {"type": "boolean"}, - "supports_none_reasoning_effort": {"type": "boolean"}, - "supports_xhigh_reasoning_effort": {"type": "boolean"}, - "supports_max_reasoning_effort": {"type": "boolean"}, - "reasoning_effort_levels": { - "type": "array", - "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, - }, - "default_reasoning_effort": { - "type": "string", - "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], - }, - "supports_adaptive_thinking": {"type": "boolean"}, - "supports_anthropic_thinking_payload": {"type": "boolean"}, - "supports_legacy_thinking": {"type": "boolean"}, - "thinking_always_on": {"type": "boolean"}, - "supports_mid_conversation_system": {"type": "boolean"}, - "supports_sampling_params": {"type": "boolean"}, - "supports_output_config": {"type": "boolean"}, - "supports_speed": {"type": "boolean"}, - "supports_fast_mode": {"type": "boolean"}, - "supported_audio_formats": { - "type": "array", - "items": { - "type": "string", - "enum": ["mp3", "wav"], - }, - }, - "vertex_ai_audio_api": { - "type": "string", - "enum": ["lyria_predict", "lyria_interactions"], - }, - "bedrock_output_config_effort_ceiling": { - "type": "string", - "enum": ["low", "medium", "high", "max", "xhigh"], - }, - "bedrock_converse_supports_strict_tools": {"type": "boolean"}, - "supports_regex_lookaround": {"type": "boolean"}, - "tpm": {"type": "number"}, - "supported_endpoints": { - "type": "array", - "items": { - "type": "string", - "enum": [ - "/v1/responses", - "/v1/embeddings", - "/v1/chat/completions", - "/v1/completions", - "/v1/messages", - "/v1/images/generations", - "/v1/realtime", - "/v1/realtime/transcription_sessions", - "/v1/images/variations", - "/v1/images/edits", - "/v1/batch", - "/v1beta/interactions", - "/v1/audio/transcriptions", - "/v1/audio/speech", - "/v1/ocr", - "/v1/videos", - "/vertex_ai/live", - "/v1/listen", - "/v1/systemone", - "/v1/decisions", - "/v1beta/interactions", - ], - }, - }, - "supported_regions": { - "type": "array", - "items": { - "type": "string", - }, - }, - "guardrail_cost_per_unit": { - "type": "object", - "additionalProperties": {"type": "number"}, - }, - "search_context_cost_per_query": { - "type": "object", - "properties": { - "search_context_size_low": {"type": "number"}, - "search_context_size_medium": {"type": "number"}, - "search_context_size_high": {"type": "number"}, - }, - "additionalProperties": False, - }, - "web_search_billing_unit": { - "type": "string", - "enum": ["per_prompt", "per_query"], - }, - "search_free_requests": { - "type": "object", - "properties": { - "count": {"type": "integer", "minimum": 1}, - "period": {"type": "string", "enum": ["day", "month"]}, - "shared_group": {"type": "string"}, - }, - "required": ["count", "period", "shared_group"], - "additionalProperties": False, - }, - "citation_cost_per_token": {"type": "number"}, - "supported_modalities": { - "type": "array", - "items": { - "type": "string", - "enum": ["text", "audio", "image", "video"], - }, - }, - "supported_output_modalities": { - "type": "array", - "items": { - "type": "string", - "enum": ["text", "image", "audio", "code", "video"], - }, - }, - "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", - "properties": { - "hours_utc": { - "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], - }, - "windows": { - "type": "array", - "items": { - "type": "object", - "properties": { - "hours_utc": { - "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], - }, - "weekdays": { - "type": "array", - "items": {"oneOf": [{"type": "integer"}, {"type": "string"}]}, - }, - }, - "required": ["hours_utc"], - "additionalProperties": False, - }, - }, - "weekday_timezone": {"type": "string"}, - "input_cost_per_token": {"type": "number"}, - "output_cost_per_token": {"type": "number"}, - "output_cost_per_reasoning_token": {"type": "number"}, - "cache_read_input_token_cost": {"type": "number"}, - "cache_creation_input_token_cost": {"type": "number"}, - }, - "additionalProperties": False, - }, - "tiered_pricing": { - "type": "array", - "items": { - "type": "object", - "properties": { - "range": { - "type": "array", - "items": {"type": "number"}, - "minItems": 2, - "maxItems": 2, - }, - "input_cost_per_token": {"type": "number"}, - "output_cost_per_token": {"type": "number"}, - "cache_read_input_token_cost": {"type": "number"}, - "cache_creation_input_token_cost": {"type": "number"}, - "output_cost_per_reasoning_token": {"type": "number"}, - "max_results_range": { - "type": "array", - "items": {"type": "number"}, - "minItems": 2, - "maxItems": 2, - }, - "input_cost_per_query": {"type": "number"}, - }, - "additionalProperties": False, - }, - }, - }, - "additionalProperties": False, - }, - } - prod_json = os.path.join(os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json") with open(prod_json, "r") as model_prices_file: actual_json = json.load(model_prices_file) @@ -1363,7 +1389,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): actual_json.pop("fallback_generalizations", None) # reserved meta key, not a model entry # Validate schema - validate(actual_json, INTENDED_SCHEMA) + validate(actual_json, _COST_MAP_ROW_SCHEMA) # Validate cost values # Define exceptions for models that are allowed to have costs > 1 @@ -1385,6 +1411,21 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): raise AssertionError(error_message) +def test_cost_map_schema_accepts_12hr_cache_write_rates_for_batch_and_ultrafast_tiers() -> None: + row: Final = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost_above_12hr_batches": 3e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 4e-6, + "cache_creation_input_token_cost_above_12hr_ultrafast": 5e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-6, + } + + assert Draft7Validator(_COST_MAP_ROW_SCHEMA).is_valid({"synthetic-12hr-cache-write-model": row}) + + def test_max_tokens_consistency(): """ Test that max_tokens == max_output_tokens for all models. @@ -2528,41 +2569,6 @@ def test_priced_router_deployment_id_stays_out_of_wildcard_model_listing(): assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.000001 -def test_reasoning_content_preserved_in_text_completion_wrapper(): - """Ensure reasoning_content is copied from delta to text_choices.""" - chunk = ModelResponseStream( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - content="Some answer text", - role="assistant", - reasoning_content="Here's my chain of thought...", - ), - ) - ], - ) - - wrapper = TextCompletionStreamWrapper( - completion_stream=None, # Not used in convert_to_text_completion_object - model="test-model", - stream_options=None, - ) - - transformed = wrapper.convert_to_text_completion_object(chunk) - - assert "choices" in transformed - assert len(transformed["choices"]) == 1 - choice = transformed["choices"][0] - assert choice["text"] == "Some answer text" - assert choice["reasoning_content"] == "Here's my chain of thought..." - - def test_anthropic_claude_4_invoke_chat_provider_config(): """Test that the Anthropic Claude 4 Invoke chat provider config is correct.""" from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( @@ -3891,6 +3897,114 @@ class TestCallbackAsyncSyncSeparation: assert my_sync_cb not in litellm._async_input_callback +class _DropsACallbackFromAnotherThreadWhenComparedToDynamodb: + def __init__(self, dropped: Callable[..., object]) -> None: + self.dropped = dropped + self.drop_thread: threading.Thread | None = None + + def __eq__(self, other: object) -> bool: + if isinstance(other, str) and other == "dynamodb" and self.drop_thread is None: + self.drop_thread = threading.Thread( + target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),) + ) + self.drop_thread.start() + self.drop_thread.join(timeout=0.2) + return self is other + + __hash__ = object.__hash__ + + +def test_function_setup_moves_async_callbacks_while_another_thread_drops_a_callback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def dropped_callback(*args: object, **kwargs: object) -> None: + return None + + async def async_callback(*args: object, **kwargs: object) -> None: + return None + + drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToDynamodb(dropped_callback) + monkeypatch.setattr(litellm.utils, "callback_list", ["initialized"]) + for list_name in ( + "callbacks", + "input_callback", + "failure_callback", + "service_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, []) + monkeypatch.setattr(litellm, "success_callback", [dropped_callback, drops_on_compare, async_callback]) + + litellm.utils.function_setup( + "completion", + litellm.utils.Rules(), + datetime.now(), + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="callback-drop-race", + ) + assert drops_on_compare.drop_thread is not None + drops_on_compare.drop_thread.join() + + assert litellm.success_callback == [drops_on_compare] + assert litellm._async_success_callback == [async_callback] + + +class _DropsACallbackFromAnotherThreadWhenComparedToSentry: + def __init__(self, dropped: Callable[..., object]) -> None: + self.dropped = dropped + self.drop_thread: threading.Thread | None = None + self.drop_finished_during_logger_setup = False + + def __eq__(self, other: object) -> bool: + if isinstance(other, str) and other == "sentry" and self.drop_thread is None: + self.drop_thread = threading.Thread( + target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),) + ) + self.drop_thread.start() + self.drop_thread.join(timeout=5) + self.drop_finished_during_logger_setup = not self.drop_thread.is_alive() + return self is other + + __hash__ = object.__hash__ + + +def test_function_setup_lets_other_threads_drop_callbacks_while_it_sets_up_loggers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def dropped_callback(*args: object, **kwargs: object) -> None: + return None + + drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToSentry(dropped_callback) + monkeypatch.setattr(litellm.utils, "callback_list", []) + for list_name in ( + "callbacks", + "input_callback", + "failure_callback", + "service_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, []) + monkeypatch.setattr(litellm, "success_callback", [drops_on_compare]) + + litellm.utils.function_setup( + "completion", + litellm.utils.Rules(), + datetime.now(), + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="logger-setup-outside-lock", + ) + assert drops_on_compare.drop_thread is not None + drops_on_compare.drop_thread.join() + + assert drops_on_compare.drop_finished_during_logger_setup + + class TestMetadataNoneHandling: """ Test that metadata=None in kwargs doesn't cause TypeError. diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py index bfaac44a483..302a7d60cdd 100644 --- a/tests/unit/test_utils_get_model_info.py +++ b/tests/unit/test_utils_get_model_info.py @@ -61,6 +61,30 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch): assert model_info["input_cost_per_token"] == 0.0 +def test_get_model_info_passes_through_gemini_tiered_audio_rates() -> None: + model: Final = "gemini-tiered-audio-rates-test" + prices: Final = { + "input_cost_per_audio_token_flex": 1e-6, + "input_cost_per_audio_token_priority": 2e-6, + "cache_read_input_audio_token_cost_flex": 3e-7, + "cache_read_input_audio_token_cost_priority": 4e-7, + "cache_read_input_audio_token_cost_batches": 5e-7, + } + litellm.register_model( + { + model: { + "litellm_provider": "gemini", + "mode": "chat", + **prices, + } + } + ) + + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + assert {key: model_info[key] for key in prices} == prices + + def test_get_model_info_ollama_chat(): from litellm.llms.ollama.completion.transformation import OllamaConfig @@ -136,6 +160,7 @@ def _enforce_bedrock_converse_models( info.get("litellm_provider") == "bedrock" and info.get("mode") == "chat" and model not in whitelist_models + and BedrockModelInfo.get_bedrock_route(model) != "claude_platform" and not ( (base_model := BedrockModelInfo.get_base_model(model)) != model and model_cost.get(base_model, {}).get("litellm_provider") == "bedrock_converse" diff --git a/tests/unit/test_uuid_helper.py b/tests/unit/test_uuid_helper.py index 589d3f96a17..f636942ba88 100644 --- a/tests/unit/test_uuid_helper.py +++ b/tests/unit/test_uuid_helper.py @@ -1,14 +1,20 @@ -import importlib +from typing import Final +from uuid import NAMESPACE_DNS, UUID, uuid5 as stdlib_uuid5 + +from litellm._uuid import uuid, uuid4 -def test_uses_fastuuid_and_uuid4_works(): - mod = importlib.import_module("litellm._uuid") - fastuuid_mod = importlib.import_module("fastuuid") - assert hasattr(mod, "uuid4") - assert hasattr(mod, "uuid") - assert mod.uuid is fastuuid_mod +def test_uuid4_returns_canonical_random_uuid() -> None: + for value in (uuid4(), uuid.uuid4()): + encoded: Final = str(value) + parsed: Final = UUID(encoded) + assert parsed.version == 4 + assert encoded == str(parsed) + assert value.bytes == parsed.bytes - # Ensure uuid4 returns something that looks like a UUID string - val = str(mod.uuid4()) - assert isinstance(val, str) - assert len(val) >= 8 + +def test_uuid5_accepts_byte_names_for_stable_ids() -> None: + namespace: Final = uuid.UUID(str(NAMESPACE_DNS)) + name: Final = "guardrail-\u00e9" + generated: Final = uuid.uuid5(namespace, name.encode("utf-8")) + assert str(generated) == str(stdlib_uuid5(NAMESPACE_DNS, name)) diff --git a/tests/unit/text_completions/__init__.py b/tests/unit/text_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py new file mode 100644 index 00000000000..b7edffc0753 --- /dev/null +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -0,0 +1,144 @@ +from collections.abc import AsyncIterator, Iterator +from typing import Final + +import pytest + +from litellm.exceptions import APIConnectionError +from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper +from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, Usage + + +def test_reasoning_content_preserved_in_text_completion_wrapper(): + """Ensure reasoning_content is copied from delta to text_choices.""" + chunk = ModelResponseStream( + id="test-id", + created=1234567890, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="Some answer text", + role="assistant", + reasoning_content="Here's my chain of thought...", + ), + ) + ], + ) + + wrapper = TextCompletionStreamWrapper( + completion_stream=None, # Not used in convert_to_text_completion_object + model="test-model", + stream_options=None, + ) + + transformed = wrapper.convert_to_text_completion_object(chunk) + + assert "choices" in transformed + assert len(transformed["choices"]) == 1 + choice = transformed["choices"][0] + assert choice["text"] == "Some answer text" + assert choice["reasoning_content"] == "Here's my chain of thought..." + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_stream", (False, True)) +@pytest.mark.parametrize("empty_choices", (False, True)) +@pytest.mark.parametrize("stream_options", (None, {}, {"include_usage": False}, {"include_usage": True})) +@pytest.mark.parametrize("usage", (None, Usage(prompt_tokens=7, completion_tokens=2, total_tokens=9), Usage())) +async def test_text_completion_stream_usage_is_opt_in( + async_stream: bool, empty_choices: bool, stream_options: dict[str, bool] | None, usage: Usage | None +) -> None: + chunk: Final = ModelResponseStream( + id="test-id", + created=1234567890, + model="test-model", + choices=[] if empty_choices else [StreamingChoices(index=0, delta=Delta(content="Hello"))], + usage=usage, + ) + + async def chunks() -> AsyncIterator[ModelResponseStream]: + yield chunk + + wrapper: Final = TextCompletionStreamWrapper( + completion_stream=chunks() if async_stream else iter((chunk,)), + model="test-model", + stream_options=stream_options, + ) + responses: Final = [response async for response in wrapper] if async_stream else list(wrapper) + expected_usage: Final = ( + {"usage": usage.model_dump(exclude_none=True)} + if stream_options and stream_options.get("include_usage") is True and usage is not None + else {} + ) + assert [response.model_dump(exclude_none=True) for response in responses] == [ + { + "id": "test-id", + "object": "text_completion", + "created": 1234567890, + "model": "test-model", + "choices": [] if empty_choices else [{"index": 0, "text": "Hello"}], + **expected_usage, + } + ] + + +def test_text_completion_conversion_rejects_non_streaming_choices() -> None: + chunk: Final = Choices() + wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model") + + with pytest.raises(Exception, match="Error occurred converting to text completion object") as error: + wrapper.convert_to_text_completion_object(chunk) + + assert str(chunk) in str(error.value) + + +@pytest.mark.parametrize("chunk", (None, "None")) +def test_sync_text_completion_stream_rejects_invalid_chunks(chunk: str | None) -> None: + wrapper: Final = TextCompletionStreamWrapper( + completion_stream=iter((chunk,)), model="test-model", custom_llm_provider="openai" + ) + + with pytest.raises(APIConnectionError, match="OpenAIException") as error: + next(wrapper) + + assert (error.value.model, error.value.llm_provider) == ("test-model", "openai") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunk", (None, "None")) +async def test_async_text_completion_stream_rejects_invalid_chunks(chunk: str | None) -> None: + async def chunks() -> AsyncIterator[str | None]: + yield chunk + + wrapper: Final = TextCompletionStreamWrapper(completion_stream=chunks(), model="test-model") + + with pytest.raises(Exception, match=r"^$") as error: + await anext(wrapper) + + assert not isinstance(error.value, StopAsyncIteration) + + +def test_sync_text_completion_stream_maps_upstream_failure_after_content() -> None: + chunk: Final = ModelResponseStream( + id="test-id", + created=1234567890, + model="test-model", + choices=[StreamingChoices(index=0, delta=Delta(content="Hello"))], + ) + + def chunks() -> Iterator[ModelResponseStream]: + yield chunk + raise RuntimeError("upstream disconnected") + + wrapper: Final = TextCompletionStreamWrapper( + completion_stream=chunks(), model="test-model", custom_llm_provider="openai" + ) + assert next(wrapper).choices[0].text == "Hello" + + with pytest.raises(APIConnectionError, match="upstream disconnected") as error: + next(wrapper) + + assert (error.value.model, error.value.llm_provider) == ("test-model", "openai") diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 788d405f99f..740342f54d9 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -94,6 +94,7 @@ CONNECTION_NAMES: Final = ( "github_copilot_host", "configurable_clientside_auth_params", "use_xai_oauth", + "chatgpt_auth_file", "fireworks_forward_user_id", "aws_batch_role_arn", "aws_bedrock_control_plane_endpoint", @@ -254,6 +255,7 @@ INTERNAL_STATE_NAMES: Final = ( "_context_compaction_state", *AGENTIC_LOOP_STATE_NAMES, "_router_weights", + "_encrypted_content_affinity_pinned", "fallback_depth", "max_fallbacks", "attempted_targets", @@ -263,6 +265,7 @@ INTERNAL_STATE_NAMES: Final = ( "litellm_trusted_callback_vars", "_litellm_addressed_response_id", "_litellm_strip_stream_usage", + "_litellm_stream_usage_injected", "client_side_timeout", "model_file_id_mapping", "acompletion", diff --git a/tests/unit/types/test_router.py b/tests/unit/types/test_router.py index 2892a829eab..7855eba0042 100644 --- a/tests/unit/types/test_router.py +++ b/tests/unit/types/test_router.py @@ -16,8 +16,8 @@ from litellm.types.router import ( holds_secret_pointer, reject_server_owned_wif_params, server_owned_wif_fields_named, - server_owned_wif_fields_written, server_owned_wif_fields_present, + server_owned_wif_fields_written, ) from litellm.types.utils import ( CustomPricingLiteLLMParams, @@ -367,6 +367,7 @@ def test_server_owned_registry_includes_anthropic_openai_and_oauth_token_exchang + openai_wif_litellm_params + oauth_token_exchange_litellm_params + github_copilot_oauth_litellm_params + + ("chatgpt_auth_file",) ) assert set(openai_wif_litellm_params) == { "openai_identity_provider_id", diff --git a/ui/Dockerfile b/ui/Dockerfile index 4b6b0c1581a..d78a155b874 100644 --- a/ui/Dockerfile +++ b/ui/Dockerfile @@ -3,7 +3,7 @@ # UI container — Next.js static export served by nginx. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 -ARG NGINX_VERSION=1.31.5-alpine3.24@sha256:34f40471dea485273c5e2a04dd5e97a682332ceb4a9adecd67de450dcb2fb390 +ARG NGINX_VERSION=1.31.8-alpine3.24@sha256:1ef7ca94d7879df82421428ce13bd1eaa412404953859cf7a5b6512c25e67a60 # ---------- builder ---------- FROM ${UI_BUILD_IMAGE} AS builder @@ -27,6 +27,8 @@ RUN npm run build # ---------- runtime ---------- FROM nginx:${NGINX_VERSION} AS runtime +RUN apk upgrade --no-cache brotli-libs + # Drop the upstream default :80 server; we own the config. RUN rm -f /etc/nginx/conf.d/default.conf diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 4802aa057c6..32777b43504 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1098,14 +1098,6 @@ "count": 1 } }, - "src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - } - }, "src/app/(dashboard)/usage/_components/components/UsagePageView.tsx": { "react-hooks/purity": { "count": 1 @@ -1767,16 +1759,16 @@ "count": 1 }, "max-params": { - "count": 15 + "count": 14 }, "no-nested-ternary": { "count": 5 }, "no-restricted-syntax": { - "count": 140 + "count": 139 }, "prefer-const": { - "count": 29 + "count": 28 } }, "src/components/object_permissions_view.tsx": { @@ -1798,9 +1790,6 @@ }, "prefer-const": { "count": 2 - }, - "react-hooks/set-state-in-effect": { - "count": 4 } }, "src/components/organization/organization_view.tsx": { @@ -2186,11 +2175,6 @@ "count": 1 } }, - "src/utils/dataUtils.test.ts": { - "max-nested-callbacks": { - "count": 1 - } - }, "tests/CreateKeyPage.expiredToken.test.tsx": { "@typescript-eslint/no-require-imports": { "count": 3 diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 944455721cd..8704dc55c38 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -28,7 +28,7 @@ "jwt-decode": "4.0.0", "lucide-react": "0.513.0", "moment": "2.31.0", - "next": "16.3.8", + "next": "16.4.1", "next-themes": "^0.4.6", "nuqs": "^2.9.4", "openai": "6.49.0", @@ -54,7 +54,7 @@ "tailwind-merge": "3.4.0", "usehooks-ts": "3.1.1", "uuid": "14.0.0", - "zod": "4.6.5" + "zod": "4.6.6" }, "devDependencies": { "@eslint/js": "9.39.2", @@ -69,8 +69,8 @@ "@types/react-copy-to-clipboard": "5.0.7", "@types/react-dom": "19.2.4", "@types/react-syntax-highlighter": "15.5.13", - "@vitest/coverage-v8": "4.1.11", - "@vitest/ui": "4.1.11", + "@vitest/coverage-v8": "4.1.12", + "@vitest/ui": "4.1.12", "eslint": "9.39.2", "eslint-config-next": "16.3.3", "eslint-config-prettier": "10.1.8", @@ -87,7 +87,7 @@ "typescript": "5.9.3", "typescript-eslint": "8.60.1", "vite": "7.3.7", - "vitest": "4.1.11" + "vitest": "4.1.12" }, "engines": { "node": ">=24.14.1", @@ -2153,9 +2153,9 @@ } }, "node_modules/@next/env": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/env/-/env-16.3.8.tgz", - "integrity": "sha512-Al9zqHVV7TJv0eFuOU4U7Lvv74PTih4Ch63sk2xCIpSTkE3udFnaOcnzP2lQVymiL7yS9Cj2iClUXlR3EQ5sEw==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/env/-/env-16.4.1.tgz", + "integrity": "sha512-5sUUrrSvs/2/2LnLoesybMGx12EtPB46fdZ6RhAw0XDGSPxU2LFyQBB7yuEiOFFJPKmBY6+q2jTr1/AMbnp0/g==", "license": "MIT" }, "node_modules/@next/eslint-plugin-next": { @@ -2170,9 +2170,9 @@ } }, "node_modules/@next/swc-darwin-arm64": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.3.8.tgz", - "integrity": "sha512-2JPRMh2nmQG5CiL7cXGL9AGwnPWJQ//cTtAUCT+w511QHk79SYz3LGv/pc5X643B/WEO0rvu3Yww0hqwt3kgeA==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.4.1.tgz", + "integrity": "sha512-U+hQ5Gj3fpzf4fveTIF3dpUFYhaSW8vHhdAnbiIomzXzg7boreYFgBqprepXFtW4sw+aTttiNmg0gew/iFNhbg==", "cpu": [ "arm64" ], @@ -2186,9 +2186,9 @@ } }, "node_modules/@next/swc-darwin-x64": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.3.8.tgz", - "integrity": "sha512-GZtCCOBKJ4leVIT/Th0llWKhD1ca92lzbQiS5R5ON9QkoiFnilFsebDae1JU2a3HWoKMEmEZWGs1AGLavVM72Q==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.4.1.tgz", + "integrity": "sha512-/hqC3V1XLUjhMASf/WridKdO9FXoEgN3HBWmtKbj00MPgnltutfraWJ3bYTpCRIgEIpbjJyOVLpYgAhh/BkdcQ==", "cpu": [ "x64" ], @@ -2202,9 +2202,9 @@ } }, "node_modules/@next/swc-linux-arm64-gnu": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.3.8.tgz", - "integrity": "sha512-O659ygeQYqneJ1fBKMpFxIFqYkYswu8IAS1OCKK/4f3ZgJJm1dRz4fVJZRi/kLLWjnBKnebOePA4WNv+sV1pVA==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.4.1.tgz", + "integrity": "sha512-GHzncHbcaS9g95kcvTNtKQW3aCLoR+xc3Jm7qrkpk1R+D2BjMHIPyvQGRPKCdMX95BsH/421jBUK06q9tEwzjw==", "cpu": [ "arm64" ], @@ -2221,9 +2221,9 @@ } }, "node_modules/@next/swc-linux-arm64-musl": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.3.8.tgz", - "integrity": "sha512-dSjKSyWpzxoO1d3DIZZcP4XJcNaKeLmxQMFOiYl5vuBRMmweIqnAhty8tAmRsvTss779cK1FtYnDMj40e4TQlg==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.4.1.tgz", + "integrity": "sha512-1SuJyDxW/PR9dHJat/ZRgcYLhuSZ+1psP4DlKb2tKIuXbjygqpLrJWRSq+onZQBv8LcmU67EIrmb6c3XTaJLjw==", "cpu": [ "arm64" ], @@ -2240,9 +2240,9 @@ } }, "node_modules/@next/swc-linux-x64-gnu": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.3.8.tgz", - "integrity": "sha512-lbqOuz3RPRcv+o9msNsJw5x4+Y1ZwPTs6vmL6DCf7i0fZfvng/F59wyeDwqHIvV0mK//RBy/jJkZ+nCKsSMXjQ==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.4.1.tgz", + "integrity": "sha512-Istng1hMGp9TGO4lUqeKLd70MPN/OOkucVTP2Nmj+gLkwBgmPqUw8CJYBCC7nbl/B/23bBCcjCI5Hsh6hazC1A==", "cpu": [ "x64" ], @@ -2259,9 +2259,9 @@ } }, "node_modules/@next/swc-linux-x64-musl": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.3.8.tgz", - "integrity": "sha512-+316WswI8ScVgZeUd+1KGaXkHhaYQzCjvH/05TZSpJ8zBizb1a4G7DtO7F12jcBIqMOtsz9ji1t48fmKtzqsGA==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.4.1.tgz", + "integrity": "sha512-aSC9wMNiZRRFVpPkPWBzRwLQxIGm/CiQcXyFF1XNcgodb+uMxdk7aSSGGpDXCbnvZdf4itjHxpOd5h7i7zaiYw==", "cpu": [ "x64" ], @@ -2278,9 +2278,9 @@ } }, "node_modules/@next/swc-win32-arm64-msvc": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.3.8.tgz", - "integrity": "sha512-ji0gd4kMYUxO+1fJBIbiBVRCjzG/lloiyCccnlebvb1ZJ5qXCPZqYg4Jl1DrrixnWNMKylzgpmMWx0yNDYXlzw==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.4.1.tgz", + "integrity": "sha512-t/myvwFDlypMUjRvbBf8qLPyuaB0iXYYdBtwMt6DM6f+60wtGr4bSYholdPqNeV5VUXZhrpMwfza829kmhS3iA==", "cpu": [ "arm64" ], @@ -2294,9 +2294,9 @@ } }, "node_modules/@next/swc-win32-x64-msvc": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.3.8.tgz", - "integrity": "sha512-WcTlaKt/TWkh5kUjdJcUmB1XgZ+1c6fz4Y9fDHL73YNSdGaUWjceeWrrlwF0nv19iABYWC4iAq1oX1w4Bn0vfg==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.4.1.tgz", + "integrity": "sha512-6ICqPo9awCkrtYNPhwkx1uU7rlFWrZRrIboDHS2LU+jv088UBYQtmJ25h6pCp27/+FZjF6BQ4Te4ErugyMeuBg==", "cpu": [ "x64" ], @@ -4490,14 +4490,14 @@ ] }, "node_modules/@vitest/coverage-v8": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-4.1.11.tgz", - "integrity": "sha512-8MVGEFnJIcdGjcbfKmeq8z0pZHH0JlVtoVZH9Q/qwUp6wyFnEJUBMrw9DCaj+ra3vShGmhavjalMIhPNxZAUcw==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-4.1.12.tgz", + "integrity": "sha512-6S7/iVknKglLAxMFxQGNppBI0iqJS2NgM6FBxavitj7KaDq/vywOWGf6dGu7F2L1Qy0G6S3uTp7aGshHrzK8JA==", "dev": true, "license": "MIT", "dependencies": { "@bcoe/v8-coverage": "^1.0.2", - "@vitest/utils": "4.1.11", + "@vitest/utils": "4.1.12", "ast-v8-to-istanbul": "^1.0.0", "istanbul-lib-coverage": "^3.2.2", "istanbul-lib-report": "^3.0.1", @@ -4511,8 +4511,8 @@ "url": "https://opencollective.com/vitest" }, "peerDependencies": { - "@vitest/browser": "4.1.11", - "vitest": "4.1.11" + "@vitest/browser": "4.1.12", + "vitest": "4.1.12" }, "peerDependenciesMeta": { "@vitest/browser": { @@ -4521,16 +4521,16 @@ } }, "node_modules/@vitest/expect": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-4.1.11.tgz", - "integrity": "sha512-VX2x5vNJXET47KAFzwERI+KRMtTTCSWTfSMKsW7JsUsXV4psq++e3DvZpuTDOpHcxytiDs6p2nhVb2tVDiiUYw==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-4.1.12.tgz", + "integrity": "sha512-vFlTt+YEu9WBpNjdGayNmlMVzrWPi1+sv853/7t2FU7hpO5gAFBKHzQmkCq5Xjpr+evegF2994LQT3hF3mRlqQ==", "dev": true, "license": "MIT", "dependencies": { "@standard-schema/spec": "^1.1.0", "@types/chai": "^5.2.2", - "@vitest/spy": "4.1.11", - "@vitest/utils": "4.1.11", + "@vitest/spy": "4.1.12", + "@vitest/utils": "4.1.12", "chai": "^6.2.2", "tinyrainbow": "^3.1.0" }, @@ -4539,13 +4539,13 @@ } }, "node_modules/@vitest/mocker": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-4.1.11.tgz", - "integrity": "sha512-2XJVD55d1o5AZous5CCGKS74g/riOj9odEt2bQpCVZeblHyHdnMeFl4jl0XjU21stf4mbjUkew2eXQZt65g5CQ==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-4.1.12.tgz", + "integrity": "sha512-fKShNpKZzupWZuY0z7egydvL0sqW4INEZBEExMBMAa6izlENEVHwoXgdaTCaBZ5fwlr4xf2o1wXC83n9RpO+NA==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/spy": "4.1.11", + "@vitest/spy": "4.1.12", "estree-walker": "^3.0.3", "magic-string": "^0.30.21" }, @@ -4566,9 +4566,9 @@ } }, "node_modules/@vitest/pretty-format": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-4.1.11.tgz", - "integrity": "sha512-yiZzPbGTS9Sr/JpFl8zHrcIkAofNbFV6k21vIgQN/cY/oxZeXhJv5sc/MBJ5jFKWmWs+oJHw0UXLZjmf931+Vw==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-4.1.12.tgz", + "integrity": "sha512-XW9O8+Mxx4VpCtkYdMJLZQQwsVryP9wk7cPKDxtXFfwSak6yRh5+2fep/Y0XF/ElBvJHq8bcYdi6aeuai0opZA==", "dev": true, "license": "MIT", "dependencies": { @@ -4579,13 +4579,13 @@ } }, "node_modules/@vitest/runner": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-4.1.11.tgz", - "integrity": "sha512-LztvUgdwMNJMIkj3hQnnxiC2Xy1zNxq928W/xhjCLaNCzqTZOudjwbQf6v9IntZGPw132i2Lq2rgTRZHD3JHNw==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-4.1.12.tgz", + "integrity": "sha512-NKlVKCr5QnCe8jijdMnmkpHHU+WTqiLFR6G0pHji9uRjXK9xJl1k1nTw6F3pOOj2PDcsDt2GKQMI7wIsaT491w==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "4.1.11", + "@vitest/utils": "4.1.12", "pathe": "^2.0.3" }, "funding": { @@ -4593,14 +4593,14 @@ } }, "node_modules/@vitest/snapshot": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-4.1.11.tgz", - "integrity": "sha512-pN7ikn1ON7h8ee4gIAp4AzyK+zBtJPzVbqOgu5LCEh4VaJVbPQcgYQYJIMGQPXVeJJq1fnfazis7a5pFNPahog==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-4.1.12.tgz", + "integrity": "sha512-ko1y1Qgb/slzQ9mr+b89Cuqud8WG1TKtsPD8r7ZUBdvVOGt85Af2HwCqHujrhwJq95hs/9RM7fcAPg57o1vlzw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "4.1.11", - "@vitest/utils": "4.1.11", + "@vitest/pretty-format": "4.1.12", + "@vitest/utils": "4.1.12", "magic-string": "^0.30.21", "pathe": "^2.0.3" }, @@ -4609,9 +4609,9 @@ } }, "node_modules/@vitest/spy": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-4.1.11.tgz", - "integrity": "sha512-apNa/prQy2qCeywhnixOHPRCgGNhvg7T4Dapfl1GahLp/R+uhBm5cPyFoNVyqsNd2h1nJxL6BqqdIjiABL60YA==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-4.1.12.tgz", + "integrity": "sha512-RGq9x4JNGNTZHEJxvN0oSiJjeSCExm0n2gZKKJZHY9cX5AVA6M7VJILBT3ZJ/tJ/0J9uWT7fETV9/btoi//vSg==", "dev": true, "license": "MIT", "funding": { @@ -4619,13 +4619,13 @@ } }, "node_modules/@vitest/ui": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-4.1.11.tgz", - "integrity": "sha512-r/rwyKoev21mWdRGSEkZOqkQ2BYy68mwjihg9M90nNRbf4NGrgzZ4cj6JNCEwlOGJkbKeMgsjlykvwKUbRr7gw==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-4.1.12.tgz", + "integrity": "sha512-eYUTHtsp52Fu6YGUw9T9173bba7YUr9j/Xj9CdhUaOb8AppLKcknHUaiZQeqJ8MiJ8C2VyKiZuUGSKdvPIKQKA==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "4.1.11", + "@vitest/utils": "4.1.12", "fflate": "^0.8.2", "flatted": "^3.4.2", "pathe": "^2.0.3", @@ -4637,17 +4637,17 @@ "url": "https://opencollective.com/vitest" }, "peerDependencies": { - "vitest": "4.1.11" + "vitest": "4.1.12" } }, "node_modules/@vitest/utils": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-4.1.11.tgz", - "integrity": "sha512-zTCVGpyFsGWBhllOyKlTw/vnr6D9qxsfSDyfbyZmTyjHw5N/VuvzHpHoQjm2ZJzn4RJgx5w4r7V0er69CmLgPQ==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-4.1.12.tgz", + "integrity": "sha512-Q3awXrwOfEV+qc3WmO5ONyE0yO7Q1gQ9/oyHxGLZ/pl7N3NO/ZztqQwKLuwMkV+CA1JYSIeO3gkHnRDIo4AVaA==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "4.1.11", + "@vitest/pretty-format": "4.1.12", "convert-source-map": "^2.0.0", "tinyrainbow": "^3.1.0" }, @@ -9710,12 +9710,12 @@ "license": "MIT" }, "node_modules/next": { - "version": "16.3.8", - "resolved": "https://registry.npmjs.org/next/-/next-16.3.8.tgz", - "integrity": "sha512-U7QEZaTini6wKrb8A8hqLLqYQyCetegKjCpJOyxk642vWoMoU1x5PyZCJFvgYgiptA8xc5j/9xYlZFO7w9Sjmw==", + "version": "16.4.1", + "resolved": "https://registry.npmjs.org/next/-/next-16.4.1.tgz", + "integrity": "sha512-uc657V62E34fOI6W5NtG9IzwsPcms0Z60grz2MJGtglgD+pTFnNt4Yw2PzaU622tfX4vhrDgLZzicE2YlhI20Q==", "license": "MIT", "dependencies": { - "@next/env": "16.3.8", + "@next/env": "16.4.1", "@swc/helpers": "0.5.23", "baseline-browser-mapping": "^2.9.19", "caniuse-lite": "^1.0.30001579", @@ -9729,15 +9729,15 @@ "node": ">=20.9.0" }, "optionalDependencies": { - "@next/swc-darwin-arm64": "16.3.8", - "@next/swc-darwin-x64": "16.3.8", - "@next/swc-linux-arm64-gnu": "16.3.8", - "@next/swc-linux-arm64-musl": "16.3.8", - "@next/swc-linux-x64-gnu": "16.3.8", - "@next/swc-linux-x64-musl": "16.3.8", - "@next/swc-win32-arm64-msvc": "16.3.8", - "@next/swc-win32-x64-msvc": "16.3.8", - "sharp": "^0.35.4" + "@next/swc-darwin-arm64": "16.4.1", + "@next/swc-darwin-x64": "16.4.1", + "@next/swc-linux-arm64-gnu": "16.4.1", + "@next/swc-linux-arm64-musl": "16.4.1", + "@next/swc-linux-x64-gnu": "16.4.1", + "@next/swc-linux-x64-musl": "16.4.1", + "@next/swc-win32-arm64-msvc": "16.4.1", + "@next/swc-win32-x64-msvc": "16.4.1", + "sharp": "^0.35.5" }, "peerDependencies": { "@opentelemetry/api": "^1.1.0", @@ -12512,19 +12512,19 @@ } }, "node_modules/vitest": { - "version": "4.1.11", - "resolved": "https://registry.npmjs.org/vitest/-/vitest-4.1.11.tgz", - "integrity": "sha512-fhACrNXUidIbGSBr5FlbuBkO7VWC1ZyLl0DO4CU2DrQoAPxX84Ysxs+HeGQpii5lZWV1Q4gBZTTu49mF+A6Edw==", + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/vitest/-/vitest-4.1.12.tgz", + "integrity": "sha512-5eWdz2yb3ajxWF3C8C4TjkXl7FGm3n2E/VlazALrQHzK7nYMSGFnCZto0bbPNl77AwkOZVLk0YnBAejUu0NiWA==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/expect": "4.1.11", - "@vitest/mocker": "4.1.11", - "@vitest/pretty-format": "4.1.11", - "@vitest/runner": "4.1.11", - "@vitest/snapshot": "4.1.11", - "@vitest/spy": "4.1.11", - "@vitest/utils": "4.1.11", + "@vitest/expect": "4.1.12", + "@vitest/mocker": "4.1.12", + "@vitest/pretty-format": "4.1.12", + "@vitest/runner": "4.1.12", + "@vitest/snapshot": "4.1.12", + "@vitest/spy": "4.1.12", + "@vitest/utils": "4.1.12", "es-module-lexer": "^2.0.0", "expect-type": "^1.3.0", "magic-string": "^0.30.21", @@ -12552,12 +12552,12 @@ "@edge-runtime/vm": "*", "@opentelemetry/api": "^1.9.0", "@types/node": "^20.0.0 || ^22.0.0 || >=24.0.0", - "@vitest/browser-playwright": "4.1.11", - "@vitest/browser-preview": "4.1.11", - "@vitest/browser-webdriverio": "4.1.11", - "@vitest/coverage-istanbul": "4.1.11", - "@vitest/coverage-v8": "4.1.11", - "@vitest/ui": "4.1.11", + "@vitest/browser-playwright": "4.1.12", + "@vitest/browser-preview": "4.1.12", + "@vitest/browser-webdriverio": "^4.1.11", + "@vitest/coverage-istanbul": "4.1.12", + "@vitest/coverage-v8": "4.1.12", + "@vitest/ui": "4.1.12", "happy-dom": "*", "jsdom": "*", "vite": "^6.0.0 || ^7.0.0 || ^8.0.0" @@ -12876,9 +12876,9 @@ } }, "node_modules/zod": { - "version": "4.6.5", - "resolved": "https://registry.npmjs.org/zod/-/zod-4.6.5.tgz", - "integrity": "sha512-v5l/aFXZQeai4awLbOpSoHecE9UiMrnfx75tEXLjNonXVARxQ5mOeipTjROUchszUNCqnE+hqAMujRsRHsut2Q==", + "version": "4.6.6", + "resolved": "https://registry.npmjs.org/zod/-/zod-4.6.6.tgz", + "integrity": "sha512-4UmzQUF5fxvfcYHlbjBmejU+w+hbDnWFx4aS8kBx7nPBHcO2SqPbpAG8snbdg+Xse4pRefR98hgjFgIMCI/UhA==", "license": "MIT", "funding": { "url": "https://github.com/sponsors/colinhacks" diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index f2a1a5e4e7f..b1b8fde44ee 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -45,7 +45,7 @@ "jwt-decode": "4.0.0", "lucide-react": "0.513.0", "moment": "2.31.0", - "next": "16.3.8", + "next": "16.4.1", "next-themes": "^0.4.6", "nuqs": "^2.9.4", "openai": "6.49.0", @@ -71,7 +71,7 @@ "tailwind-merge": "3.4.0", "usehooks-ts": "3.1.1", "uuid": "14.0.0", - "zod": "4.6.5" + "zod": "4.6.6" }, "devDependencies": { "@eslint/js": "9.39.2", @@ -86,8 +86,8 @@ "@types/react-copy-to-clipboard": "5.0.7", "@types/react-dom": "19.2.4", "@types/react-syntax-highlighter": "15.5.13", - "@vitest/coverage-v8": "4.1.11", - "@vitest/ui": "4.1.11", + "@vitest/coverage-v8": "4.1.12", + "@vitest/ui": "4.1.12", "eslint": "9.39.2", "eslint-config-next": "16.3.3", "eslint-config-prettier": "10.1.8", @@ -104,9 +104,13 @@ "typescript": "5.9.3", "typescript-eslint": "8.60.1", "vite": "7.3.7", - "vitest": "4.1.11" + "vitest": "4.1.12" }, "overrides": { + "@litellm/lens-ui": { + "next": "$next" + }, + "zod": "$zod", "prismjs": "1.30.0", "js-yaml": "4.3.2", "brace-expansion": "5.0.12", diff --git a/ui/litellm-dashboard/public/assets/aria-agent.png b/ui/litellm-dashboard/public/assets/aria-agent.png new file mode 100644 index 00000000000..870773933e3 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/aria-agent.png differ diff --git a/ui/litellm-dashboard/public/assets/logos/claude.svg b/ui/litellm-dashboard/public/assets/logos/claude.svg new file mode 100644 index 00000000000..ede30652b25 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/claude.svg @@ -0,0 +1,27 @@ + +Claude diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx index ae898645de4..b4fb7586e6d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx @@ -134,13 +134,16 @@ describe("BudgetTable", () => { } }); - it("should not make the reset column sortable", async () => { + it("should request duration sorting from the server", async () => { const user = userEvent.setup(); - renderWithProviders(); + const onSortingChange = vi.fn(); + renderWithProviders(); await showColumn(user, "budget_duration"); const headers = screen.getAllByRole("columnheader").map((header) => header.textContent); expect(headers).toContain("Reset"); - expect(screen.queryByTestId("sort-header-budget_duration")).not.toBeInTheDocument(); + await user.click(screen.getByTestId("sort-header-budget_duration")); + expect(onSortingChange).toHaveBeenCalled(); + expect(onSortingChange.mock.calls[0][0]([])).toEqual([{ id: "budget_duration", desc: false }]); }); it("should ask the list for a new sort when a sortable header is clicked", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx index fb894322208..3ec291c1158 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx @@ -139,8 +139,7 @@ export const getBudgetTableColumns = ({ accessorKey: "budget_duration", filterFn: serverFilter, meta: { title: "Reset" }, - // "7d"/"30d" sort lexicographically, not chronologically, so the route does not offer it. - enableSorting: false, + sortDescFirst: false, header: ({ column }) => , size: 110, cell: ({ row }) => , diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx new file mode 100644 index 00000000000..0d74b714ee6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx @@ -0,0 +1,78 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { GuardrailDeleteModal } from "./GuardrailDeleteModal"; +import { deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils"; + +vi.mock("@/components/networking", () => ({ deleteGuardrailCall: vi.fn() })); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + +const GUARDRAIL = { + guardrail_id: "gr-1", + guardrail_name: "word-block", + litellm_params: { guardrail: "litellm_content_filter", mode: "pre_call", default_on: true }, +}; + +const renderModal = () => { + const onClose = vi.fn(); + const onDeleted = vi.fn(); + renderWithProviders( + , + ); + return { onClose, onDeleted }; +}; + +const confirmDelete = () => { + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { + target: { value: GUARDRAIL.guardrail_name }, + }); + fireEvent.click(screen.getByRole("button", { name: "Delete" })); +}; + +describe("GuardrailDeleteModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("warns about name references and only enables Delete once the exact name is typed", () => { + renderModal(); + + const warning = screen.getByText(/silently stop running it/); + expect(warning).toHaveTextContent("Keys, teams, and policies that list this guardrail"); + expect(warning).toHaveTextContent("Policy pipeline steps that use it will block every request with a 400 error"); + expect(warning).toHaveTextContent("unless the step's on_error is next or allow"); + const confirmButton = screen.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { target: { value: "word" } }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { + target: { value: GUARDRAIL.guardrail_name }, + }); + expect(confirmButton).toBeEnabled(); + }); + + it("shows a success toast naming the guardrail and calls onDeleted after a successful delete", async () => { + vi.mocked(deleteGuardrailCall).mockResolvedValue({}); + const { onClose, onDeleted } = renderModal(); + + confirmDelete(); + + await waitFor(() => expect(onDeleted).toHaveBeenCalledTimes(1)); + expect(deleteGuardrailCall).toHaveBeenCalledWith("tok", GUARDRAIL.guardrail_id); + expect(toast.success).toHaveBeenCalledWith('Guardrail "word-block" deleted successfully'); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(onClose).toHaveBeenCalledTimes(1); + }); + + it("shows an error toast and does not call onDeleted when the delete fails", async () => { + vi.mocked(deleteGuardrailCall).mockRejectedValue(new Error("boom")); + const { onClose, onDeleted } = renderModal(); + + confirmDelete(); + + await waitFor(() => expect(onClose).toHaveBeenCalledTimes(1)); + expect(toast.fromError).toHaveBeenCalledWith("Failed to delete guardrail"); + expect(toast.success).not.toHaveBeenCalled(); + expect(onDeleted).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx new file mode 100644 index 00000000000..a3ed15b630f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx @@ -0,0 +1,68 @@ +import React, { useState } from "react"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import { deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; + +const DANGLING_REFERENCE_WARNING = + "Keys, teams, and policies that list this guardrail will silently stop running it. Policy pipeline steps that use it will block every request with a 400 error, unless the step's on_error is next or allow (or its on_fail, when on_error is not set)."; + +export interface DeletableGuardrail { + guardrail_id: string; + guardrail_name?: string | null; + litellm_params?: { + guardrail?: string; + mode?: unknown; + default_on?: boolean | null; + } | null; +} + +interface GuardrailDeleteModalProps { + guardrail: DeletableGuardrail | null; + accessToken: string | null; + onClose: () => void; + onDeleted: () => void; +} + +export function GuardrailDeleteModal({ guardrail, accessToken, onClose, onDeleted }: GuardrailDeleteModalProps) { + const [isDeleting, setIsDeleting] = useState(false); + const guardrailName = guardrail ? guardrail.guardrail_name || guardrail.guardrail_id : ""; + const providerDisplayName = getGuardrailLogoAndName(guardrail?.litellm_params?.guardrail || "").displayName; + + const handleDelete = async () => { + if (!guardrail || !accessToken) return; + setIsDeleting(true); + try { + await deleteGuardrailCall(accessToken, guardrail.guardrail_id); + toast.success(`Guardrail "${guardrailName}" deleted successfully`); + onClose(); + onDeleted(); + } catch { + toast.fromError("Failed to delete guardrail"); + onClose(); + } finally { + setIsDeleting(false); + } + }; + + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx new file mode 100644 index 00000000000..46bc09e7383 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx @@ -0,0 +1,51 @@ +import { Code, Trash2 } from "lucide-react"; +import React, { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { DeletableGuardrail, GuardrailDeleteModal } from "./GuardrailDeleteModal"; + +interface GuardrailSettingsActionsProps { + guardrail: DeletableGuardrail & { guardrail_definition_location?: string }; + isEditing: boolean; + onEdit: () => void; + onEditCode: () => void; + accessToken: string | null; + onDeleted: () => void; +} + +export function GuardrailSettingsActions({ + guardrail, + isEditing, + onEdit, + onEditCode, + accessToken, + onDeleted, +}: GuardrailSettingsActionsProps) { + const [isModalOpen, setIsModalOpen] = useState(false); + + if (isEditing || guardrail.guardrail_definition_location === "config") return null; + + return ( + <> + {guardrail.litellm_params?.guardrail === "custom_code" ? ( + + ) : ( + + )} + + setIsModalOpen(false)} + onDeleted={onDeleted} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx new file mode 100644 index 00000000000..19eb838da6b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx @@ -0,0 +1,135 @@ +import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import GuardrailsPanel from "./GuardrailsPanel"; +import { + deleteGuardrailCall, + getGuardrailInfo, + getGuardrailProviderSpecificParams, + getGuardrailsList, + getGuardrailUISettings, +} from "@/components/networking"; +import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils"; + +vi.mock("@/components/networking", () => ({ + getGuardrailsList: vi.fn(), + deleteGuardrailCall: vi.fn(), + getGuardrailInfo: vi.fn(), + getGuardrailUISettings: vi.fn(), + getGuardrailProviderSpecificParams: vi.fn(), + updateGuardrailCall: vi.fn(), +})); + +vi.mock("./add_guardrail_form", () => ({ __esModule: true, default: () => null })); +vi.mock("./guardrail_garden", () => ({ __esModule: true, default: () => null })); +vi.mock("./GuardrailTestPlayground", () => ({ __esModule: true, default: () => null })); +vi.mock("./TeamGuardrailsTab", () => ({ TeamGuardrailsTab: () => null })); + +const DB_GUARDRAIL = { + guardrail_id: "db-guardrail-1", + guardrail_name: "db-word-block", + litellm_params: { guardrail: "litellm_content_filter", mode: "pre_call", default_on: true }, + guardrail_info: null, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "db", +}; + +const CONFIG_GUARDRAIL = { + ...DB_GUARDRAIL, + guardrail_id: "config-guardrail-1", + guardrail_name: "cfg-word-block", + guardrail_definition_location: "config", +}; + +beforeAll(() => { + Object.defineProperty(window, "matchMedia", { + writable: true, + value: vi.fn().mockImplementation((query: string) => ({ + matches: false, + media: query, + onchange: null, + addListener: vi.fn(), + removeListener: vi.fn(), + addEventListener: vi.fn(), + removeEventListener: vi.fn(), + dispatchEvent: vi.fn(), + })), + }); +}); + +const UI_SETTINGS = { + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], +}; + +describe("GuardrailsPanel delete from the guardrail detail Settings card", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [DB_GUARDRAIL, CONFIG_GUARDRAIL] }); + vi.mocked(getGuardrailInfo).mockImplementation(async (_token: string, id: string) => + id === CONFIG_GUARDRAIL.guardrail_id ? CONFIG_GUARDRAIL : DB_GUARDRAIL, + ); + vi.mocked(getGuardrailUISettings).mockResolvedValue(UI_SETTINGS); + vi.mocked(getGuardrailProviderSpecificParams).mockResolvedValue({}); + }); + + const renderDetail = (guardrailId: string) => + renderWithProviders(, { + searchParams: `?guardrail=${guardrailId}`, + }); + + it("deletes a DB guardrail only after the admin types its name, then returns to the refreshed table", async () => { + vi.mocked(deleteGuardrailCall).mockResolvedValue({}); + renderDetail(DB_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + + const modal = within(await screen.findByRole("dialog")); + expect( + modal.getByText(/pipeline steps that use it will block every request with a 400 error/i), + ).toBeInTheDocument(); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + + const confirmationInput = modal.getByPlaceholderText(DB_GUARDRAIL.guardrail_name); + fireEvent.change(confirmationInput, { target: { value: "db-word" } }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(confirmationInput, { target: { value: DB_GUARDRAIL.guardrail_name } }); + expect(confirmButton).toBeEnabled(); + + fireEvent.click(confirmButton); + + await waitFor(() => expect(deleteGuardrailCall).toHaveBeenCalledWith("test-token", DB_GUARDRAIL.guardrail_id)); + expect(await screen.findByRole("table")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + expect(getGuardrailsList).toHaveBeenCalledTimes(2); + }); + + it("keeps the admin on the detail page when the delete request fails", async () => { + vi.mocked(deleteGuardrailCall).mockRejectedValue(new Error("boom")); + renderDetail(DB_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + const modal = within(await screen.findByRole("dialog")); + fireEvent.change(modal.getByPlaceholderText(DB_GUARDRAIL.guardrail_name), { + target: { value: DB_GUARDRAIL.guardrail_name }, + }); + fireEvent.click(modal.getByRole("button", { name: "Delete" })); + + await waitFor(() => expect(deleteGuardrailCall).toHaveBeenCalledTimes(1)); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.getByRole("button", { name: "Delete" })).toBeInTheDocument(); + expect(screen.queryByRole("table")).not.toBeInTheDocument(); + }); + + it("does not offer Delete on a config guardrail's detail page", async () => { + renderDetail(CONFIG_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + expect(await screen.findByRole("img", { name: "Config guardrail details" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Delete" })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx index 7be120a4cb7..ac824f63be9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx @@ -2,6 +2,7 @@ import { type UrlUpdateEvent } from "nuqs/adapters/testing"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import GuardrailsPanel from "./GuardrailsPanel"; import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils"; vi.mock("@/components/networking", () => ({ @@ -9,6 +10,8 @@ vi.mock("@/components/networking", () => ({ deleteGuardrailCall: vi.fn(), })); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + vi.mock("./add_guardrail_form", () => ({ __esModule: true, default: () =>

Mock Add Guardrail Form
, @@ -140,13 +143,36 @@ describe("GuardrailsPanel", () => { expect(modal.getByText("Delete Guardrail")).toBeInTheDocument(); expect(modal.getByText("test-guardrail-1")).toBeInTheDocument(); expect(modal.getByText("Test Provider")).toBeInTheDocument(); + expect( + modal.getByText(/pipeline steps that use it will block every request with a 400 error/i), + ).toBeInTheDocument(); - fireEvent.click(modal.getByRole("button", { name: "Delete" })); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(confirmButton); await waitFor(() => { expect(mockDeleteGuardrailCall).toHaveBeenCalledWith("test-token", "test-guardrail-1"); }); - expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2); + await waitFor(() => expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2)); + expect(toast.success).toHaveBeenCalledWith('Guardrail "Test Guardrail" deleted successfully'); + }); + + it("should show an error toast and not refetch when the row delete fails", async () => { + mockDeleteGuardrailCall.mockRejectedValue(new Error("boom")); + renderWithProviders(); + fireEvent.click(screen.getByText("Guardrails")); + + fireEvent.click(await screen.findByTestId("delete-button")); + const modal = within(await screen.findByRole("dialog")); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(modal.getByRole("button", { name: "Delete" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(toast.fromError).toHaveBeenCalledWith("Failed to delete guardrail"); + expect(toast.success).not.toHaveBeenCalled(); + expect(mockGetGuardrailsList).toHaveBeenCalledTimes(1); }); it("should mount every tab panel up front so panel state survives tab switches", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index 901e39004f3..b53d016e950 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -2,7 +2,7 @@ import { parseAsString, useQueryState } from "nuqs"; import React, { useState, useEffect } from "react"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { ChevronDown, Code, Plus } from "lucide-react"; -import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking"; +import { getGuardrailsList } from "@/components/networking"; import { buttonVariants } from "@/components/ui/button"; import { DropdownMenu, @@ -16,10 +16,8 @@ import GuardrailTable from "./guardrail_table"; import { isAdminRole } from "@/utils/roles"; import GuardrailInfoView from "./guardrail_info"; import GuardrailTestPlayground from "./GuardrailTestPlayground"; -import { toast } from "@/lib/toast"; import { Guardrail } from "@/components/guardrails/types"; -import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { GuardrailDeleteModal } from "./GuardrailDeleteModal"; import { CustomCodeModal } from "./custom_code"; import GuardrailGarden from "./guardrail_garden"; import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; @@ -38,9 +36,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isCustomCodeModalVisible, setIsCustomCodeModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(false); - const [isDeleting, setIsDeleting] = useState(false); const [guardrailToDelete, setGuardrailToDelete] = useState(null); - const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [selectedGuardrailId, setSelectedGuardrailId] = useQueryState( "guardrail", parseAsString.withOptions({ history: "push" }), @@ -97,40 +93,10 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole fetchGuardrails(); }; - const handleDeleteClick = (guardrailId: string, guardrailName: string) => { - const guardrail = guardrailsList.find((g) => g.guardrail_id === guardrailId) || null; - setGuardrailToDelete(guardrail); - setIsDeleteModalOpen(true); + const handleDeleteClick = (guardrailId: string) => { + setGuardrailToDelete(guardrailsList.find((g) => g.guardrail_id === guardrailId) || null); }; - const handleDeleteConfirm = async () => { - if (!guardrailToDelete || !accessToken) return; - - setIsDeleting(true); - try { - await deleteGuardrailCall(accessToken, guardrailToDelete.guardrail_id); - toast.success(`Guardrail "${guardrailToDelete.guardrail_name}" deleted successfully`); - await fetchGuardrails(); - } catch (error) { - console.error("Error deleting guardrail:", error); - toast.fromError("Failed to delete guardrail"); - } finally { - setIsDeleting(false); - setIsDeleteModalOpen(false); - setGuardrailToDelete(null); - } - }; - - const handleDeleteCancel = () => { - setIsDeleteModalOpen(false); - setGuardrailToDelete(null); - }; - - const providerDisplayName = - guardrailToDelete && guardrailToDelete.litellm_params - ? getGuardrailLogoAndName(guardrailToDelete.litellm_params.guardrail).displayName - : undefined; - return (
@@ -184,6 +150,10 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole { + closeGuardrailDetail(); + fetchGuardrails(); + }} accessToken={accessToken} isAdmin={isAdmin} /> @@ -210,24 +180,11 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole onSuccess={handleSuccess} /> - setGuardrailToDelete(null)} + onDeleted={fetchGuardrails} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx index 1689cf93289..732c50647b1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx @@ -84,7 +84,8 @@ const saveChanges = async (user: ReturnType) => { await user.click(screen.getByText("Save Changes")); }; -const renderView = () => render( {}} accessToken="123" isAdmin />); +const renderView = () => + render( {}} onDeleted={() => {}} accessToken="123" isAdmin />); const lastPayload = () => vi.mocked(networking.updateGuardrailCall).mock.calls.at(-1)?.[2]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 836921104ff..ae057689b37 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -10,6 +10,7 @@ vi.mock("@/components/networking", () => ({ getGuardrailUISettings: vi.fn(), getGuardrailProviderSpecificParams: vi.fn(), updateGuardrailCall: vi.fn(), + deleteGuardrailCall: vi.fn(), })); // Mock ContentFilterManager @@ -65,7 +66,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); // Wait for the loading to complete and data to be rendered await waitFor(() => { @@ -103,7 +106,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); expect(await screen.findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); }); @@ -131,7 +136,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); const logo = await screen.findByAltText("Presidio PII logo"); expect(logo).toHaveAttribute("src", expect.stringContaining("microsoft_azure.svg")); @@ -162,7 +169,7 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); const { container } = render( - {}} accessToken="123" isAdmin={true} />, + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, ); await waitFor(() => { @@ -212,7 +219,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); await waitFor(() => { expect(screen.getByText("PII Entity Configuration")).toBeInTheDocument(); @@ -245,7 +254,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); vi.mocked(networking.updateGuardrailCall).mockResolvedValue({ status: "success" }); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); await waitFor(() => { expect(screen.getByText("Settings")).toBeInTheDocument(); @@ -331,7 +342,9 @@ describe("Guardrail Info", () => { }); vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); expect(await screen.findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); @@ -355,10 +368,90 @@ describe("Guardrail Info when the guardrail cannot be loaded", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); const onClose = vi.fn(); - render(); + render( + {}} + accessToken="123" + isAdmin={true} + />, + ); expect(await screen.findByText("Guardrail not found")).toBeInTheDocument(); fireEvent.click(screen.getByRole("button", { name: /back to guardrails/i })); expect(onClose).toHaveBeenCalledTimes(1); }); + + describe("Delete from the Settings card", () => { + const dbGuardrail = { + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { guardrail: "presidio", mode: "pre_call", default_on: true }, + guardrail_definition_location: "db", + }; + const uiSettings = { + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call"], + }; + const mockDbGuardrail = () => { + vi.mocked(networking.getGuardrailInfo).mockResolvedValue(dbGuardrail); + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue(uiSettings); + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + }; + + const openDeleteModalAndConfirm = async () => { + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + const modal = within(await screen.findByRole("dialog")); + expect(modal.getByText(/silently stop running it/)).toBeInTheDocument(); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(confirmButton); + }; + + it("deletes the guardrail and calls onDeleted once the admin types its name", async () => { + mockDbGuardrail(); + vi.mocked(networking.deleteGuardrailCall).mockResolvedValue({}); + const onDeleted = vi.fn(); + render( + {}} onDeleted={onDeleted} accessToken="tok" isAdmin />, + ); + + await openDeleteModalAndConfirm(); + + await waitFor(() => expect(onDeleted).toHaveBeenCalledTimes(1)); + expect(networking.deleteGuardrailCall).toHaveBeenCalledWith("tok", "123"); + }); + + it("hides Delete while the settings are being edited", async () => { + mockDbGuardrail(); + render( {}} onDeleted={() => {}} accessToken="tok" isAdmin />); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Edit Settings" })); + + expect(await screen.findByRole("button", { name: "Save Changes" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Delete" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + }); + + it("does not call onDeleted when the delete request fails", async () => { + mockDbGuardrail(); + vi.mocked(networking.deleteGuardrailCall).mockRejectedValue(new Error("boom")); + const onDeleted = vi.fn(); + render( + {}} onDeleted={onDeleted} accessToken="tok" isAdmin />, + ); + + await openDeleteModalAndConfirm(); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(networking.deleteGuardrailCall).toHaveBeenCalledTimes(1); + expect(onDeleted).not.toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index 2cfe2b79886..eacacf5e30e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -33,6 +33,7 @@ import { } from "./GuardrailFormField"; import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal"; +import { GuardrailSettingsActions } from "./GuardrailSettingsActions"; import { GuardrailModeCard } from "./GuardrailModeDisplay"; import { GuardrailReadOnlyDetails } from "./GuardrailReadOnlyDetails"; import { @@ -72,11 +73,12 @@ const SectionHeading: React.FC<{ children: React.ReactNode }> = ({ children }) = export interface GuardrailInfoProps { guardrailId: string; onClose: () => void; + onDeleted: () => void; accessToken: string | null; isAdmin: boolean; } -const GuardrailInfoView: React.FC = ({ guardrailId, onClose, accessToken, isAdmin }) => { +const GuardrailInfoView: React.FC = ({ guardrailId, onClose, onDeleted, accessToken, isAdmin }) => { const [guardrailData, setGuardrailData] = useState(null); const [guardrailProviderSpecificParams, setGuardrailProviderSpecificParams] = useState(null); const [loading, setLoading] = useState(true); @@ -686,25 +688,21 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, {isAdmin && ( -
-

Guardrail Settings

+
+

Guardrail Settings

{isConfigGuardrail && ( )} - {!isEditing && - !isConfigGuardrail && - (guardrailData.litellm_params?.guardrail === "custom_code" ? ( - - ) : ( - - ))} + setIsEditing(true)} + onEditCode={() => setCustomCodeModalVisible(true)} + accessToken={accessToken} + onDeleted={onDeleted} + />
{isEditing ? ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx index b957393be64..1891a580088 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx @@ -49,6 +49,9 @@ const requestUrl = (input: RequestInfo | URL): string => { }; const fakeFetch = (input: RequestInfo | URL) => { + if (requestUrl(input).endsWith("/models")) { + return Promise.resolve(Response.json({ data: [{ id: "claude-sonnet-5-5" }] })); + } if (requestUrl(input).includes("/model_group/info")) { return Promise.resolve(Response.json({ data: [{ model_group: "claude-sonnet-5-5", mode: "chat" }] })); } @@ -141,7 +144,7 @@ describe("HomePage", () => { expect(screen.queryByRole("link", { name: /create new virtual key/i })).not.toBeInTheDocument(); }); - it("offers LiteAdmin on Home inside the dashboard frame", async () => { + it("should offer Aria Agent on Home inside the dashboard frame", async () => { render( @@ -149,7 +152,7 @@ describe("HomePage", () => { , ); - const entry = screen.getByRole("region", { name: "Ask LiteAdmin" }); + const entry = screen.getByRole("region", { name: "Ask Aria Agent" }); expect(await within(entry).findByPlaceholderText("What would you like to do?")).toBeInTheDocument(); expect(within(entry).getByRole("button", { name: "Send message" })).toBeDisabled(); expect(await screen.findByText("$501.74")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index 24f76065f34..744db87ad11 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -121,6 +121,7 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment { /** False for config.yaml-defined deployments, which the update and delete routes refuse. */ db_model?: boolean | null; base_model?: string | null; + blocked?: boolean | null; mode?: string | null; created_at?: string | null; updated_at?: string | null; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx index a66f31a0dc0..54decbfff3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx @@ -102,6 +102,7 @@ const AccessGroupBudgetModal: React.FC = ({ id={id} value={value || null} onChange={(next) => onChange(next ?? undefined)} + placeholder="Keep current" /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx index 6f7e04f92c6..2d1c24685f0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx @@ -106,10 +106,7 @@ export function AutoRoutersPanel({ > Add Auto Router - - Choose a classifier to route each request to a model. Called like any other model, so clients keep using a - single model name. - + Choose models for your auto router. { expect(await screen.findByLabelText(/Max Budget/)).toHaveValue(2.5); }); + it("keeps the stored reset period when Keep current is selected", async () => { + GET.mockResolvedValue({ + data: { + access_groups: [ + { + ...BUDGETED_GROUP, + budget: { ...BUDGETED_GROUP.budget, budget_duration: "24h" }, + }, + ], + }, + }); + renderPanel(); + await openActions("premium"); + await userEvent.click(await screen.findByText("Edit budget")); + + await userEvent.click(await screen.findByRole("combobox", { name: "Reset Budget" })); + const keepCurrentOption = await screen.findByRole("option", { name: "Keep current" }); + expect(keepCurrentOption).toBeInTheDocument(); + await userEvent.click(keepCurrentOption); + await userEvent.click(screen.getByRole("button", { name: "Save Budget" })); + + await waitFor(() => + expect(PUT).toHaveBeenCalledWith("/access_group/{access_group}/budget", { + params: { path: { access_group: "premium" } }, + body: { max_budget: 2.5 }, + }), + ); + expect(PUT.mock.calls[0][1].body).not.toHaveProperty("budget_duration"); + }); + it("refuses to save a budget with every field blank", async () => { renderPanel(); await openActions("shared"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/PassThroughPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/PassThroughPanel.tsx index 7ae7e976651..ad08ff48a5b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/PassThroughPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/PassThroughPanel.tsx @@ -4,8 +4,6 @@ import PassThroughSettings from "@/components/PassThroughSettings/PassThroughSet import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; export default function PassThroughPanel() { - const { accessToken, userRole, userId: userID, premiumUser } = useAuthorized(); - return ( - - ); + const { accessToken, userRole, userId: userID } = useAuthorized(); + return ; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx index 839eb406200..724a72a77f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx @@ -66,6 +66,85 @@ const baseProps = { }; describe("OrganizationsTable", () => { + it.each(["max_budget", "tpm_limit", "rpm_limit"])( + "sorts %s numerically with Unlimited as infinity", + async (field) => { + const user = userEvent.setup(); + const organizations = [20, 3, null, 100, 0].map((value, index) => + makeOrganization({ + organization_id: `org-${index}`, + organization_alias: `Value ${value}`, + litellm_budget_table: { [field]: value }, + }), + ); + renderWithProviders(, { + searchParams: { sort_by: field, sort_order: "asc", page_size: "25" }, + }); + const names = () => + screen + .getAllByRole("row") + .slice(1) + .map((row) => within(row).getByText(/^Value /).textContent); + expect(names()).toEqual(["Value 0", "Value 3", "Value 20", "Value 100", "Value null"]); + if (field === "max_budget") { + await user.click(screen.getByTestId("sort-header-max_budget")); + } else { + await user.click(screen.getByRole("button", { name: "Sort options for TPM or RPM" })); + await user.click( + await screen.findByRole("menuitem", { name: `${field === "tpm_limit" ? "TPM" : "RPM"} descending` }), + ); + } + await waitFor(() => expect(names()).toEqual(["Value null", "Value 100", "Value 20", "Value 3", "Value 0"])); + }, + ); + + it.each(["asc", "desc"])("groups both All Proxy Models representations when sorting %s", (direction) => { + const organizations = [ + makeOrganization({ organization_id: "org-empty", organization_alias: "Name Empty", models: [] }), + makeOrganization({ organization_id: "org-z", organization_alias: "Name Z", models: ["z"] }), + makeOrganization({ + organization_id: "org-wildcard", + organization_alias: "Name Wildcard", + models: ["all-proxy-models"], + }), + makeOrganization({ organization_id: "org-a", organization_alias: "Name A", models: ["a"] }), + ]; + renderWithProviders(, { + searchParams: { sort_by: "models", sort_order: direction }, + }); + const names = screen + .getAllByRole("row") + .slice(1) + .map((row) => within(row).getByText(/^Name /).textContent); + const models = direction === "asc" ? ["Name A", "Name Z"] : ["Name Z", "Name A"]; + expect(names).toEqual([...models, "Name Empty", "Name Wildcard"]); + expect(screen.getAllByText("All Proxy Models")).toHaveLength(2); + }); + + it.each(["models", "members"])("sorts %s by its displayed data", async (field) => { + const user = userEvent.setup(); + const organizations = [ + { organization_id: "org-zed", organization_alias: "Zed", models: ["z"], members: [] }, + { + organization_id: "org-ace", + organization_alias: "Ace", + models: ["a"], + members: [{ user_id: "1" }, { user_id: "2" }], + }, + { + organization_id: "org-mid", + organization_alias: "Mid", + models: ["m"], + members: [{ user_id: "3" }], + }, + ].map(makeOrganization); + renderWithProviders(); + await user.click(screen.getByTestId(`sort-header-${field}`)); + await waitFor(() => expect(bodyRowAliases()).toEqual(["Ace", "Mid", "Zed"])); + await user.click(screen.getByTestId(`sort-header-${field}`)); + await waitFor(() => expect(bodyRowAliases()).toEqual(["Zed", "Mid", "Ace"])); + }); + it("renders every column header", () => { renderWithProviders(); for (const header of [ @@ -75,7 +154,6 @@ describe("OrganizationsTable", () => { "Spend (USD)", "Budget (USD)", "Models", - "TPM / RPM Limits", "Members", ]) { expect(screen.getByText(header)).toBeInTheDocument(); @@ -87,7 +165,7 @@ describe("OrganizationsTable", () => { for (const header of ["Spend (USD)", "Budget (USD)", "Members"]) { expect(screen.getByRole("columnheader", { name: header })).toHaveClass("text-right"); } - for (const header of ["Organization Name", "TPM / RPM Limits"]) { + for (const header of ["Organization Name"]) { expect(screen.getByRole("columnheader", { name: header })).not.toHaveClass("text-right"); } }); @@ -270,7 +348,7 @@ describe("OrganizationsTable URL state", () => { it("falls back to sorting by creation date for a ?sort_by= column that cannot be sorted", () => { renderWithProviders(, { - searchParams: "?sort_by=members&sort_order=asc", + searchParams: "?sort_by=actions&sort_order=asc", }); expect(bodyRowAliases()).toEqual(["Zed", "Mid", "Ace"]); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx index a9ac0b7e699..df1c475ebb1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx @@ -62,6 +62,7 @@ const OrganizationsTable: React.FC = ({ onPaginationChange={onPaginationChange} columns={columns} getRowId={(organization, index) => organization.organization_id || String(index)} + defaultColumnVisibility={{ tpm_limit: false, rpm_limit: false }} sortingMode="client" sorting={sorting} onSortingChange={onSortingChange} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx index 5f170a32941..dcbf141314b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx @@ -3,7 +3,7 @@ import { ColumnDef } from "@tanstack/react-table"; import { MoreHorizontal, Pencil, Trash2 } from "lucide-react"; -import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DataTableMultiSortHeader, DataTableSortHeader } from "@/components/shared/DataTable"; import { DateCell, IdentityCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; import { Organization } from "@/components/networking"; import { buttonVariants } from "@/components/ui/button"; @@ -137,38 +137,67 @@ export const getOrganizationsTableColumns = ({ }, { id: "max_budget", + accessorFn: (organization) => getOrganizationBudget(organization).max_budget ?? Number.POSITIVE_INFINITY, + sortUndefined: "last", + sortDescFirst: true, meta: { title: "Budget (USD)", numeric: true }, - header: "Budget (USD)", + header: ({ column }) => , size: 120, - enableSorting: false, + enableSorting: true, cell: ({ row }) => ( ), }, { id: "models", + accessorFn: (organization) => + organization.models?.length && !organization.models.includes("all-proxy-models") + ? [...organization.models].sort().join(", ") + : undefined, + sortUndefined: "last", + sortDescFirst: false, meta: { title: "Models", skeleton: "chips" }, - header: "Models", + header: ({ column }) => , size: 260, - enableSorting: false, + enableSorting: true, cell: ({ row }) => , }, { id: "limits", meta: { title: "TPM / RPM Limits" }, - header: "TPM / RPM Limits", + header: ({ table }) => ( + + ), size: 150, enableSorting: false, cell: ({ row }) => , }, { id: "members", + accessorFn: (organization) => organization.members?.length ?? 0, + sortUndefined: "last", + sortDescFirst: true, meta: { title: "Members", numeric: true }, - header: "Members", + header: ({ column }) => , size: 100, - enableSorting: false, + enableSorting: true, cell: ({ row }) => {row.original.members?.length ?? 0} Members, }, + ...(["tpm_limit", "rpm_limit"] as const).map( + (field): ColumnDef => ({ + id: field, + accessorFn: (organization) => getOrganizationBudget(organization)[field] ?? Number.POSITIVE_INFINITY, + sortUndefined: "last", + sortDescFirst: true, + enableHiding: false, + }), + ), { id: "actions", meta: { className: "text-right", headerClassName: "text-right" }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/useOrganizationsTableState.ts b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/useOrganizationsTableState.ts index 20a54a25ba1..ffe67c0aa14 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/useOrganizationsTableState.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/useOrganizationsTableState.ts @@ -4,7 +4,17 @@ const FILTER_COLUMNS = ["org_id"] as const; type FilterColumn = (typeof FILTER_COLUMNS)[number]; const TABLE_STATE_OPTIONS: UrlTableStateOptions = { - sortFields: ["organization_id", "organization_alias", "created_at", "spend"], + sortFields: [ + "organization_id", + "organization_alias", + "created_at", + "spend", + "max_budget", + "models", + "members", + "tpm_limit", + "rpm_limit", + ], defaultSort: { id: "created_at", desc: true }, defaultPageSize: 25, filterColumns: FILTER_COLUMNS, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.test.tsx index da2e4debb3c..4903a431ff7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.test.tsx @@ -1,4 +1,5 @@ import { fireEvent, render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; import { ChatComposer } from "./ChatComposer"; @@ -111,4 +112,23 @@ describe("ChatComposer", () => { expect(screen.getByTestId("chat-composer-input")).toHaveFocus(); }); + + it.each([ + { label: "Find a model", control: }, + { + label: "Choose a model", + control: ( + + ), + }, + ])("should preserve focus when clicking the $label toolbar control", async ({ label, control }) => { + renderComposer({ tools: control }); + + await userEvent.setup().click(screen.getByLabelText(label)); + + expect(screen.getByLabelText(label)).toHaveFocus(); + expect(screen.getByPlaceholderText("Send a message")).not.toHaveFocus(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx index 5d2edbdfc94..c33b5e992fc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx @@ -51,6 +51,14 @@ export function ChatComposer({ } }; + const handleToolbarClick = (event: React.MouseEvent) => { + const control = (event.target as HTMLElement).closest( + "button, input, select, textarea, a[href], [contenteditable='true'], [tabindex]", + ); + if (event.defaultPrevented || (control && event.currentTarget.contains(control))) return; + event.currentTarget.parentElement?.querySelector("[data-slot=input-group-control]")?.focus(); + }; + return (
{showSuggestions && suggestions.length > 0 && ( @@ -99,12 +107,7 @@ export function ChatComposer({ { - if ((event.target as HTMLElement).closest("button")) { - return; - } - event.currentTarget.parentElement?.querySelector("[data-slot=input-group-control]")?.focus(); - }} + onClick={handleToolbarClick} >
{tools}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx index 35434830d98..d7dfb8151fc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx @@ -80,6 +80,42 @@ describe("GeneralSettings General tab", () => { vi.mocked(deleteConfigFieldSetting).mockClear(); }); + it.each([ + ["maximum_spend_logs_retention_period", null, "7d"], + ["maximum_health_check_retention_period", "30d", "14d"], + ["maximum_daily_tag_spend_retention_period", undefined, "90d"], + ["maximum_spend_logs_cleanup_run_budget", "5m", "10m"], + ["maximum_spend_logs_cleanup_batch_timeout", "", "45s"], + ["mcp_client_id_header", "x-client-id", "X-MCP-Client.ID"], + ])("displays, updates and resets the String setting %s", async (fieldName, initialValue, updatedValue) => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + ...SETTINGS_FIXTURE, + { + field_name: fieldName, + field_type: "String", + field_value: initialValue, + field_description: "string setting", + stored_in_db: initialValue != null, + field_default_value: null, + }, + ]); + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("tab", { name: "General" })); + const input = await screen.findByRole("textbox", { name: fieldName }); + expect(input).toHaveValue(initialValue ?? ""); + fireEvent.change(input, { target: { value: updatedValue } }); + const row = screen.getByRole("row", { name: new RegExp(fieldName) }); + await user.click(within(row).getByRole("button", { name: "Update" })); + expect(vi.mocked(updateConfigFieldSetting).mock.calls).toEqual([["token", fieldName, updatedValue]]); + expect(within(row).getByText("In DB")).toBeInTheDocument(); + const resetIcon = row.querySelectorAll("td")[3].querySelector("svg"); + expect(resetIcon).not.toBeNull(); + await user.click(resetIcon as unknown as Element); + expect(vi.mocked(deleteConfigFieldSetting).mock.calls).toEqual([["token", fieldName]]); + expect(input).toHaveValue(""); + }); + it("updates max_ui_session_budget with its own value, not the value at its filtered index", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 5d28adb7ef3..67712d78a13 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -68,6 +68,15 @@ const SettingValueEditor: React.FC<{ /> ); } + if (setting.field_type === "String") { + return ( + onChange(setting.field_name, event.target.value)} + /> + ); + } if (setting.field_type === "Integer") { return ( { }); it("should handle search errors and show notification", async () => { - const error = new Error("Search failed"); + const error = new Error( + JSON.stringify({ + error: { + message: "User not allowed to access search tool: exa-search-2", + type: "user_search_tool_access_denied", + code: "403", + }, + }), + ); vi.mocked(networking.searchToolQueryCall).mockRejectedValue(error); const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); const user = userEvent.setup(); @@ -281,11 +289,29 @@ describe("SearchToolTester", () => { const searchButton = screen.getByRole("button", { name: /search/i }); await user.click(searchButton); await waitFor(() => { - expect(toast.fromError).toHaveBeenCalledWith("Failed to query search tool"); + expect(toast.fromError).toHaveBeenCalledWith(error); }); consoleSpy.mockRestore(); }); + it("should not show empty results or add history when search fails", async () => { + vi.mocked(networking.searchToolQueryCall).mockRejectedValue(new Error("Search failed")); + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + const user = userEvent.setup(); + render(); + const input = screen.getByPlaceholderText("Enter your search query..."); + fireEvent.change(input, { target: { value: "test query" } }); + await user.click(screen.getByRole("button", { name: /search/i })); + + await waitFor(() => { + expect(toast.fromError).toHaveBeenCalled(); + }); + expect(screen.queryByText("No results found")).not.toBeInTheDocument(); + expect(screen.queryByText("test query")).not.toBeInTheDocument(); + expect(screen.getByText("Test your search tool")).toBeInTheDocument(); + consoleSpy.mockRestore(); + }); + it("should maintain search history", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx index 7dcef3fc6ba..646120443c7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx @@ -60,7 +60,7 @@ export const SearchToolTester: React.FC = ({ searchToolNa setSearchHistory((prev) => [historyEntry, ...prev]); } catch (error) { console.error("Error querying search tool:", error); - toast.fromError("Failed to query search tool"); + toast.fromError(error); } finally { setIsLoading(false); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsagePanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsagePanel.integration.test.tsx new file mode 100644 index 00000000000..39627329b85 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsagePanel.integration.test.tsx @@ -0,0 +1,411 @@ +import { useState, type ComponentProps } from "react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { uiSpendLogsCall, userListCall, dailyActivityExportCall, type UserInfo } from "@/components/networking"; +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import { + chooseSelectOption, + renderWithProviders, + screen, + testQueryClient, + waitFor, + within, +} from "../../../../../../tests/test-utils"; +import ClientUsagePanel from "./ClientUsagePanel"; +import { UsageClient, type AgentRow } from "./overview/agentCatalog"; + +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + uiSpendLogsCall: vi.fn(), + dailyActivityExportCall: vi.fn(), + userListCall: vi.fn(), +})); + +const desktopTag = "User-Agent: claude-cli/2.1.263 (external, claude-desktop-3p)"; +const agent: AgentRow = { + id: "claude-desktop", + label: "Claude Desktop", + description: "Desktop assistant", + kind: "agent", + logo: null, + spend: 12.5, + tokens: 9000, + requests: 3, + users: 2, + tags: [desktopTag], +}; +const metrics: SpendMetrics = { + spend: 12.5, + total_tokens: 9000, + prompt_tokens: 8000, + completion_tokens: 1000, + cache_read_input_tokens: 5000, + cache_creation_input_tokens: 1000, + api_requests: 3, + successful_requests: 3, + failed_requests: 0, +}; +const daily: DailyData[] = [ + { + date: "2026-10-22", + metrics, + breakdown: { + models: {}, + model_groups: {}, + providers: {}, + api_keys: {}, + mcp_servers: {}, + entities: { [desktopTag]: { metrics, metadata: {}, api_key_breakdown: {} } }, + }, + }, +]; +const defaultProps: ComponentProps = { + client: UsageClient.ClaudeDesktop, + userId: null, + canChooseUser: true, + onUserChange: vi.fn(), + onClose: vi.fn(), + accessToken: "test-token", + startTime: new Date("2026-10-22T00:00:00Z"), + endTime: new Date("2026-10-22T23:59:59.999Z"), + dateValue: { from: new Date("2026-10-22T00:00:00Z"), to: new Date("2026-10-22T23:59:59.999Z") }, + presetShortLabel: null, + onDateChange: vi.fn(), + agent, + daily: [], + loading: false, +}; +const personLog = { + request_id: "request-1", + user: "person-1", + startTime: "2026-10-22T12:00:00Z", + model: "test-model", + spend: 3, + total_tokens: 1000, + prompt_tokens: 800, + completion_tokens: 200, + status: "success", + request_tags: [desktopTag], + metadata: {}, +}; +const person: UserInfo = { + user_id: "person-1", + user_email: "person@example.com", + user_alias: null, + allowed_cache_controls: [], + key_count: 0, + model_max_budget: {}, + model_spend: {}, + models: [], + policies: [], + spend: 0, + teams: [], +}; + +describe("ClientUsagePanel", () => { + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + vi.stubGlobal("requestIdleCallback", (callback: IdleRequestCallback) => { + callback({ didTimeout: false, timeRemaining: () => 50 }); + return 1; + }); + const userPage = { + users: [person], + total: 1, + page: 1, + page_size: 50, + total_pages: 1, + }; + const requestPage = { + data: [personLog], + total: 1, + page: 1, + total_pages: 1, + total_is_capped: false, + }; + const exportedRows = [ + { + date: "2026-10-22", + entity_id: desktopTag, + model: "test-model", + spend: 12.5, + flat_cost: 0, + prompt_tokens: 8000, + completion_tokens: 1000, + api_requests: 3, + successful_requests: 3, + failed_requests: 0, + cache_read_input_tokens: 5000, + cache_creation_input_tokens: 1000, + }, + ]; + vi.mocked(dailyActivityExportCall).mockResolvedValue( + Object.assign(new Blob([JSON.stringify(exportedRows)], { type: "application/json" }), { + text: async () => JSON.stringify(exportedRows), + }), + ); + vi.mocked(userListCall).mockResolvedValue(userPage); + vi.mocked(uiSpendLogsCall).mockResolvedValue(requestPage); + vi.stubGlobal( + "fetch", + vi.fn(async () => Response.json({ results: [] })), + ); + }); + + afterEach(() => vi.unstubAllGlobals()); + + it("should immediately show the clicked all-users summary and logo without requiring a person", () => { + renderWithProviders(); + + const panel = screen.getByRole("dialog", { name: "Claude Desktop" }); + expect(within(panel).getByRole("img", { name: "Claude Desktop logo" })).toBeInTheDocument(); + expect(panel).toHaveTextContent("All users"); + expect(panel).toHaveTextContent(/Spend\s*\$12\.50/); + expect(panel).toHaveTextContent(/Total tokens\s*9K/); + expect(panel).toHaveTextContent(/Requests\s*3/); + expect(screen.queryByRole("status", { name: "Loading client usage" })).not.toBeInTheDocument(); + expect(screen.queryByText(/Select a user to/)).not.toBeInTheDocument(); + expect(uiSpendLogsCall).not.toHaveBeenCalled(); + }); + + it("should filter to a person inside the open panel and keep absent cache details unavailable", async () => { + const user = userEvent.setup(); + const onUserChange = vi.fn(); + const onClose = vi.fn(); + function ControlledPanel() { + const [userId, setUserId] = useState(null); + return ( + { + setUserId(id); + onUserChange(id); + }} + /> + ); + } + renderWithProviders(); + + await chooseSelectOption(user, screen.getByRole("combobox", { name: "Filter by person" }), /person@example.com/); + expect(onUserChange).toHaveBeenCalledWith("person-1"); + expect(onClose).not.toHaveBeenCalled(); + await waitFor(() => + expect(uiSpendLogsCall).toHaveBeenCalledWith( + expect.objectContaining({ + params: expect.objectContaining({ user_id: "person-1" }), + }), + ), + ); + const panel = screen.getByRole("dialog", { name: "Claude Desktop" }); + await waitFor(() => expect(panel).toHaveTextContent(/Spend\s*\$3\.00/)); + expect(panel).toHaveTextContent("Selected user"); + expect(panel).toHaveTextContent(/Cache reads\s*Unavailable/); + expect(panel).toHaveTextContent(/Cache writes\s*Unavailable/); + expect(panel).toHaveTextContent(/Uncached input\s*Unavailable/); + expect(panel).toHaveTextContent(/Output\s*200/); + + await user.click(screen.getByRole("button", { name: "All users" })); + expect(onUserChange).toHaveBeenLastCalledWith(null); + expect(onClose).not.toHaveBeenCalled(); + expect(screen.getByRole("dialog", { name: "Claude Desktop" })).toHaveTextContent(/Spend\s*\$12\.50/); + }); + + it.each(["Back to Usage", "Close"])("should return to the overview through %s", async (button) => { + const user = userEvent.setup(); + const onClose = vi.fn(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: button })); + expect(onClose).toHaveBeenCalledOnce(); + }); + + it("should preserve summary totals and mark partial model details unavailable", async () => { + const partialAgent = { ...agent, spend: 25, tokens: 18000, requests: 6 }; + renderWithProviders(); + + const panel = screen.getByRole("dialog", { name: "Claude Desktop" }); + expect(panel).toHaveTextContent(/Spend\s*\$25\.00/); + expect(within(panel).getAllByText("Unavailable")).toHaveLength(4); + expect(await screen.findByText(/Model usage does not match these client totals yet/)).toBeInTheDocument(); + expect(screen.queryByRole("heading", { name: "Daily usage" })).not.toBeInTheDocument(); + expect(screen.queryByText("5,000")).not.toBeInTheDocument(); + }); + + it("should not expose a cached all-users summary without a permitted user scope", () => { + renderWithProviders(); + + expect(screen.getByRole("alert")).toHaveTextContent("Sign in to view your client usage"); + expect(screen.queryByText("$12.50")).not.toBeInTheDocument(); + expect(screen.queryByText("Total tokens")).not.toBeInTheDocument(); + expect(screen.queryByRole("combobox", { name: "Filter by person" })).not.toBeInTheDocument(); + expect(screen.queryByRole("heading", { name: "Daily usage" })).not.toBeInTheDocument(); + expect(uiSpendLogsCall).not.toHaveBeenCalled(); + expect(screen.queryByRole("button", { name: "Try again" })).not.toBeInTheDocument(); + }); + + it("should show a failed summary as an error and retry without closing or showing stale totals", async () => { + const user = userEvent.setup(); + const onRetry = vi.fn(); + const onClose = vi.fn(); + renderWithProviders(); + + expect(screen.getByRole("alert")).toHaveTextContent("Could not load client totals"); + expect(screen.queryByText("$12.50")).not.toBeInTheDocument(); + expect(screen.queryByText("No Claude Desktop usage recorded in this date range.")).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Try again" })); + expect(onRetry).toHaveBeenCalledOnce(); + expect(onClose).not.toHaveBeenCalled(); + }); + it("should change the panel time range through the existing picker without closing", async () => { + const user = userEvent.setup(); + const onDateChange = vi.fn(); + const onClose = vi.fn(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /22 Oct/ })); + await user.click(screen.getByRole("button", { name: /Last 7 days/ })); + await user.click(screen.getByRole("button", { name: "Apply" })); + expect(onDateChange).toHaveBeenCalledWith( + expect.objectContaining({ from: expect.any(Date), to: expect.any(Date) }), + "7d", + ); + expect(onClose).not.toHaveBeenCalled(); + expect(screen.getByRole("dialog", { name: "Claude Desktop" })).toBeInTheDocument(); + }); + + it("should show the client model breakdown using the overview controls", async () => { + renderWithProviders(); + expect(await screen.findByText("test-model")).toBeInTheDocument(); + expect(screen.getByRole("heading", { name: "Daily usage" })).toBeInTheDocument(); + expect(screen.getByRole("radiogroup", { name: "Metric" })).toBeInTheDocument(); + expect(screen.getByRole("radio", { name: "Tokens" })).toBeInTheDocument(); + expect(screen.getByRole("radio", { name: "Weekly" })).toBeInTheDocument(); + }); + + it("should show personal client usage without loading the global user list or all-user export", async () => { + renderWithProviders(); + + const panel = screen.getByRole("dialog", { name: "Claude Desktop" }); + await waitFor(() => expect(panel).toHaveTextContent(/Spend\s*\$3\.00/)); + expect(panel).toHaveTextContent("My usage"); + expect(screen.queryByRole("combobox", { name: "Filter by person" })).not.toBeInTheDocument(); + expect(userListCall).not.toHaveBeenCalled(); + expect(dailyActivityExportCall).not.toHaveBeenCalled(); + expect(uiSpendLogsCall).toHaveBeenCalledWith( + expect.objectContaining({ params: expect.objectContaining({ user_id: "person-1" }) }), + ); + expect(vi.mocked(uiSpendLogsCall).mock.lastCall?.[0].params).not.toHaveProperty("team_id"); + }); + + it("should keep team and member filters intersected inside the open panel without fetching global users", async () => { + const user = userEvent.setup(); + const onClose = vi.fn(); + const onTeamChange = vi.fn(); + const onUserChange = vi.fn(); + const teamOptions = [{ value: "engineering", label: "Engineering" }]; + const memberOptions = [ + { value: "person-1", label: "John" }, + { value: "person-2", label: "Maya" }, + ]; + const johnLog = { ...personLog, team_id: "engineering" }; + const mayaLog = { ...johnLog, request_id: "request-2", user: "person-2", spend: 7 }; + vi.mocked(uiSpendLogsCall).mockImplementation(async (request) => { + const logs = [johnLog, mayaLog].filter((log) => !request.params?.user_id || log.user === request.params.user_id); + return { data: logs, total: logs.length, page: 1, total_pages: 1, total_is_capped: false }; + }); + function TeamPanel() { + const [teamId, setTeamId] = useState(null); + const [userId, setUserId] = useState("person-1"); + const changeTeam = (id: string | null) => { + setTeamId(id); + setUserId(id ? null : "person-1"); + onTeamChange(id); + }; + return ( + { + setUserId(id); + onUserChange(id); + }} + /> + ); + } + renderWithProviders(); + + await chooseSelectOption(user, screen.getByRole("combobox", { name: "Usage scope" }), "Engineering"); + expect(onTeamChange).toHaveBeenCalledWith("engineering"); + const panel = screen.getByRole("dialog", { name: "Claude Desktop" }); + await waitFor(() => expect(panel).toHaveTextContent(/Spend\s*\$10\.00/)); + expect(panel).toHaveTextContent("Engineering · All members"); + expect(uiSpendLogsCall).toHaveBeenLastCalledWith( + expect.objectContaining({ params: expect.objectContaining({ team_id: "engineering" }) }), + ); + expect(vi.mocked(uiSpendLogsCall).mock.lastCall?.[0].params).not.toHaveProperty("user_id"); + + await user.click(screen.getByRole("combobox", { name: "Team member" })); + expect(await screen.findByRole("option", { name: "Maya" })).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "John" })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: "Outside user" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("option", { name: "Maya" })); + expect(onUserChange).toHaveBeenCalledWith("person-2"); + await waitFor(() => expect(panel).toHaveTextContent(/Spend\s*\$7\.00/)); + expect(panel).toHaveTextContent("Engineering · Maya"); + expect(uiSpendLogsCall).toHaveBeenLastCalledWith( + expect.objectContaining({ params: expect.objectContaining({ team_id: "engineering", user_id: "person-2" }) }), + ); + + await chooseSelectOption(user, screen.getByRole("combobox", { name: "Team member" }), "All members"); + expect(onUserChange).toHaveBeenLastCalledWith(null); + await waitFor(() => expect(panel).toHaveTextContent(/Spend\s*\$10\.00/)); + expect(onClose).not.toHaveBeenCalled(); + expect(userListCall).not.toHaveBeenCalled(); + expect(dailyActivityExportCall).not.toHaveBeenCalled(); + }); + + it.each([ + [null, "Choose an administered team before viewing another member's usage"], + ["engineering", "Choose a member of the selected team"], + ])("should block requests and cached summaries for an invalid scope in team %s", (teamId, scopeError) => { + renderWithProviders( + , + ); + + expect(screen.getByRole("alert")).toHaveTextContent(scopeError); + expect(screen.queryByText("$12.50")).not.toBeInTheDocument(); + expect(screen.queryByText("Total tokens")).not.toBeInTheDocument(); + expect(screen.queryByRole("heading", { name: "Daily usage" })).not.toBeInTheDocument(); + expect(uiSpendLogsCall).not.toHaveBeenCalled(); + expect(dailyActivityExportCall).not.toHaveBeenCalled(); + expect(userListCall).not.toHaveBeenCalled(); + }); + it("should recover a newer model export by refreshing the cached summary on retry", async () => { + const user = userEvent.setup(); + function StaleSummaryPanel() { + const [summary, setSummary] = useState({ ...agent, spend: 12 }); + return setSummary(agent)} />; + } + renderWithProviders(); + expect(await screen.findByText(/Model usage does not match these client totals/)).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Retry model breakdown" })); + expect(await screen.findByText("test-model")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsagePanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsagePanel.tsx new file mode 100644 index 00000000000..3ca8f2e038d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsagePanel.tsx @@ -0,0 +1,362 @@ +"use client"; + +import { useMemo } from "react"; +import { ArrowLeft, Coins } from "lucide-react"; +import { useTagSummary } from "@/app/(dashboard)/hooks/tags/useTagSummary"; +import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; +import UserDropdown from "@/components/common_components/UserDropdown"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import type { ClientUsageScopeOption } from "../hooks/clientUsageScope"; +import { stackedUsageColor } from "@/components/shared/charts"; +import { Button } from "@/components/ui/button"; +import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from "@/components/ui/sheet"; +import { Skeleton } from "@/components/ui/skeleton"; +import { EMPTY_DAILY_ACTIVITY_METADATA, toDailyData } from "@/components/UsagePage/dailyActivityApi"; +import type { DailyData } from "@/components/UsagePage/types"; +import claudeLogo from "../../../../../../public/assets/logos/claude.svg"; +import type { UsageClient } from "../hooks/clientUsageData"; +import { useClientUsage } from "../hooks/useClientUsage"; +import { topAgents, type AgentRow } from "./overview/agentCatalog"; +import { useClientAgentModels } from "../hooks/useClientAgentModels"; +import UsageBreakdownChart from "./overview/UsageBreakdownChart"; +import { clientAgentDetail, type ClientAgentDetail } from "./overview/clientAgentDetail"; +import { modelColorMap, formatCompact, formatUsd, overviewTotals } from "./overview/overviewData"; +import { reconciledTokenBreakdown } from "./overview/tokenBreakdown"; +import { Panel, Stat } from "./overview/Primitives"; +import { AgentMark } from "./overview/TopAgents"; + +interface ClientUsagePanelProps { + client: UsageClient; + userId: string | null; + canChooseUser: boolean; + teamId?: string | null; + scopeError?: string | null; + teams?: readonly ClientUsageScopeOption[]; + members?: readonly ClientUsageScopeOption[]; + defaultScopeLabel?: string; + onTeamChange?: (teamId: string | null) => void; + onUserChange: (userId: string | null) => void; + onClose: () => void; + accessToken: string | null; + startTime: Date | null; + endTime: Date | null; + dateValue: DateRangePickerValue; + presetShortLabel: string | null; + onDateChange: (value: DateRangePickerValue, preset: string | null) => void; + colorSource?: readonly DailyData[]; + agent: AgentRow | null; + daily: readonly DailyData[]; + loading: boolean; + summaryFailed?: boolean; + onRetry?: () => void; +} + +const dateLabel = (date: Date | null) => date?.toLocaleDateString("en-US", { month: "short", day: "numeric" }) ?? "—"; +const countLabel = (value: number | null) => (value === null ? "Unavailable" : value.toLocaleString()); + +export default function ClientUsagePanel({ + client, + userId, + canChooseUser, + teamId = null, + scopeError: accessError = null, + teams = [], + members = [], + defaultScopeLabel = "All users", + onTeamChange, + onUserChange, + onClose, + accessToken, + startTime, + endTime, + dateValue, + presetShortLabel, + onDateChange, + colorSource, + agent, + daily, + loading, + summaryFailed = false, + onRetry, +}: ClientUsagePanelProps) { + const label = client === "claude-desktop" ? "Claude Desktop" : "Claude Code"; + const scopedRequest = Boolean(userId || teamId) && !accessError; + const personRequest = { + accessToken, + startTime, + endTime, + userId, + teamId, + client: scopedRequest ? client : null, + }; + const personUsage = useClientUsage(personRequest); + const allUsersEnabled = canChooseUser && !scopedRequest && !accessError; + const teamLabel = teams.find((team) => team.value === teamId)?.label ?? teamId; + const personLabel = members.find((member) => member.value === userId)?.label ?? "Selected user"; + const selectedScopeLabel = (() => { + if (teamId) return `${teamLabel} · ${userId ? personLabel : "All members"}`; + if (!canChooseUser) return "My usage"; + return userId ? "Selected user" : "All users"; + })(); + const summaryQuery = useTagSummary(startTime, endTime, allUsersEnabled && agent === null); + const activeAgent = useMemo( + () => agent ?? topAgents(summaryQuery.data ?? []).find((row) => row.id === client) ?? null, + [summaryQuery.data, agent, client], + ); + const modelRequest = { accessToken, startTime, endTime, agent: activeAgent, enabled: allUsersEnabled }; + const agentModels = useClientAgentModels(modelRequest); + const agentDaily = agentModels.results.length > 0 ? agentModels.results : daily; + const agentDetail = useMemo( + () => (activeAgent ? clientAgentDetail(activeAgent, agentDaily) : null), + [activeAgent, agentDaily], + ); + const colorByModel = useMemo(() => (colorSource ? modelColorMap(colorSource) : undefined), [colorSource]); + const personTotals = overviewTotals(personUsage.data.metadata ?? EMPTY_DAILY_ACTIVITY_METADATA, null); + const pending = + !accessError && (scopedRequest ? personUsage.loading : (loading || summaryQuery.isLoading) && activeAgent === null); + const summaryError = summaryQuery.isError || summaryFailed ? "Could not load client totals" : null; + const scopeError = + accessError ?? (!canChooseUser && !scopedRequest ? "Sign in to view your client usage" : summaryError); + const error = accessError ?? (scopedRequest ? personUsage.error : scopeError); + const summary = scopedRequest ? personTotals : agentDetail?.summary; + const rawBreakdown = scopedRequest + ? { + inputTokens: personTotals.inputTokens, + outputTokens: personTotals.outputTokens, + cacheReadTokens: personUsage.cacheAvailable ? personTotals.cacheReadTokens : null, + cacheWriteTokens: personUsage.cacheAvailable ? personTotals.cacheWriteTokens : null, + } + : agentDetail?.tokenBreakdown; + const breakdown = rawBreakdown && summary ? reconciledTokenBreakdown(summary.tokens, rawBreakdown) : undefined; + const modelResults = scopedRequest ? toDailyData(personUsage.data) : agentModels.results; + const chartLoading = scopedRequest ? personUsage.loading : agentModels.loading; + const chartError = scopedRequest ? personUsage.error : agentModels.error; + + const canFilterPerson = canChooseUser && !teamId && !accessError; + const canRetry = !accessError && (scopedRequest || canChooseUser); + const ready = !pending && !error; + const retryModels = () => { + void summaryQuery.refetch(); + onRetry?.(); + agentModels.reload(); + }; + + return ( + { + if (!open) onClose(); + }} + > + + + +
+ +
+ {label} + + {dateLabel(startTime)} – {dateLabel(endTime)} · {selectedScopeLabel} + +
+
+
+
+
+ +
+ {teams.length > 0 && onTeamChange && ( +
+ + ({ ...team, value: `team:${team.value}` })), + ]} + allowClear={false} + onValueChange={(value) => onTeamChange(value?.startsWith("team:") ? value.slice(5) : null)} + /> +
+ )} + {teamId && ( +
+ + ({ ...member, value: `user:${member.value}` })), + ]} + allowClear={false} + onValueChange={(value) => onUserChange(value?.startsWith("user:") ? value.slice(5) : null)} + /> +
+ )} + {canFilterPerson && ( +
+
+ + {userId && ( + + )} +
+ +
+ )} + {pending && ( +
+ +
+ )} + {error && ( +
+

Could not load client usage

+

{error}

+ {canRetry && ( + + )} +
+ )} + {ready && !summary && ( +

No {label} usage recorded in this date range.

+ )} + {ready && summary && ( + <> +
+ + + +
+ {summary.requests === 0 ? ( +

+ No {label} requests in the selected range. Try another date range. + {(canChooseUser || teamId) && " You can also choose another person."} +

+ ) : ( + <> + {!chartError && ( + + )} + {chartError && ( +
+

+ {chartError} +

+ +
+ )} + {breakdown && } + + )} +

+ {scopedRequest + ? "Based on retained request logs within the exact selected time range." + : "All-user totals use UTC daily aggregates for the selected dates. Filtering by person uses the exact selected times. Cache values are reported counts."}{" "} + Spend is the cost recorded by LiteLLM. +

+ + )} +
+
+
+ ); +} + +function SummaryStat({ label, value }: { label: string; value: string }) { + return ( +
+ +
+ ); +} + +function TokenExplanation({ total, breakdown }: { total: number; breakdown: ClientAgentDetail["tokenBreakdown"] }) { + const { inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens } = breakdown; + const freshInput = + inputTokens !== null && cacheReadTokens !== null && cacheWriteTokens !== null + ? inputTokens - cacheReadTokens - cacheWriteTokens + : null; + const rows = [ + { label: "Cache reads", value: cacheReadTokens, color: stackedUsageColor(0) }, + { label: "Cache writes", value: cacheWriteTokens, color: stackedUsageColor(1) }, + { label: "Uncached input", value: freshInput, color: stackedUsageColor(2) }, + { label: "Output", value: outputTokens, color: stackedUsageColor(3) }, + ]; + const complete = rows.every((row) => row.value !== null); + return ( + + {complete && total > 0 && ( + + )} + {cacheReadTokens !== null && total > 0 && ( +

+ {((cacheReadTokens / total) * 100).toFixed(1)}% of these tokens are cache reads. +

+ )} +
+ {rows.map((row) => ( +
+
+ + {row.label} +
+
{countLabel(row.value)}
+
+ ))} +
+

+ Total tokens include input, cached input and output. Cached input is priced differently, so token count alone + does not explain spend. Compare the same person and dates when checking the count shown by a client. +

+
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsageSelect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsageSelect.tsx new file mode 100644 index 00000000000..d0b0c7750d6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/ClientUsageSelect.tsx @@ -0,0 +1,25 @@ +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { parseUsageClient, UsageClient } from "./overview/agentCatalog"; + +const CLIENT_OPTIONS = [ + { value: UsageClient.ClaudeDesktop, label: "Claude Desktop" }, + { value: UsageClient.ClaudeCode, label: "Claude Code" }, +]; + +export default function ClientUsageSelect({ + value, + onChange, +}: { + value: UsageClient | null; + onChange: (client: UsageClient | null) => void; +}) { + return ( + onChange(parseUsageClient(client))} + /> + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx deleted file mode 100644 index d971f5d0506..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx +++ /dev/null @@ -1,83 +0,0 @@ -import { screen } from "@testing-library/react"; -import { beforeAll, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "@/../tests/test-utils"; -import UsageAIChatPanel from "./UsageAIChatPanel"; - -beforeAll(() => { - if (typeof window !== "undefined" && !window.ResizeObserver) { - window.ResizeObserver = class ResizeObserver { - observe() {} - unobserve() {} - disconnect() {} - } as any; - } -}); - -vi.mock("@/components/networking", () => ({ - modelHubCall: vi.fn().mockResolvedValue({ - data: [{ model_group: "gpt-4" }, { model_group: "claude-3-opus" }], - }), - usageAiChatStream: vi.fn(), -})); - -const defaultProps = { - open: true, - onClose: vi.fn(), - accessToken: "test-token", -}; - -describe("UsageAIChatPanel", () => { - it("should render the panel when open", () => { - renderWithProviders(); - - expect(screen.getByText("Ask AI")).toBeInTheDocument(); - expect(screen.getByText("Ask about your spend, models, keys, and trends")).toBeInTheDocument(); - }); - - it("should render model selector", () => { - renderWithProviders(); - - // One library paints the prompt as its own text node and the other leaves it on the input's - // placeholder attribute, so either one means the user is being told what to pick. - const prompt = "Select a model (optional, defaults to gpt-4o-mini)"; - expect(screen.queryAllByText(prompt).length + screen.queryAllByPlaceholderText(prompt).length).toBeGreaterThan(0); - }); - - it("should render empty state message when no conversation", () => { - renderWithProviders(); - - expect(screen.getByText("Ask a question about your usage")).toBeInTheDocument(); - }); - - it("should render the send button", () => { - renderWithProviders(); - - expect(screen.getByText("Send")).toBeInTheDocument(); - }); - - it("should render input placeholder", () => { - renderWithProviders(); - - expect(screen.getByPlaceholderText("Ask about your usage...")).toBeInTheDocument(); - }); - - it("should render clear chat button", () => { - renderWithProviders(); - - expect(screen.getByText("Clear chat")).toBeInTheDocument(); - }); - - it("should have the panel element even when closed (just off-screen)", () => { - renderWithProviders(); - - expect(screen.getByTestId("usage-ai-chat-panel")).toBeInTheDocument(); - expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-full"); - }); - - it("should not have translate-x-full class when open", () => { - renderWithProviders(); - - expect(screen.getByTestId("usage-ai-chat-panel")).not.toHaveClass("translate-x-full"); - expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-0"); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx deleted file mode 100644 index dc4291d25e1..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx +++ /dev/null @@ -1,385 +0,0 @@ -import React, { useEffect, useRef, useState } from "react"; -import ReactMarkdown from "react-markdown"; -import { MessageSquare, Sparkles, X } from "lucide-react"; -import { Button } from "@/components/ui/button"; -import { - Combobox, - ComboboxContent, - ComboboxEmpty, - ComboboxInput, - ComboboxItem, - ComboboxList, -} from "@/components/ui/combobox"; -import { Textarea } from "@/components/ui/textarea"; -import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; -import { modelHubCall, usageAiChatStream, UsageAiToolCallEvent } from "@/components/networking"; - -interface ToolCallStep { - tool_name: string; - tool_label: string; - arguments: Record; - status: "running" | "complete" | "error"; - error?: string; -} - -interface ChatMessage { - role: "user" | "assistant"; - content: string; - toolCalls?: ToolCallStep[]; -} - -interface UsageAIChatPanelProps { - open: boolean; - onClose: () => void; - accessToken: string | null; -} - -const TOOL_ICONS: Record = { - get_usage_data: "📊", - get_team_usage_data: "👥", - get_tag_usage_data: "🏷️", -}; - -const ToolCallDisplay: React.FC<{ step: ToolCallStep }> = ({ step }) => { - const icon = TOOL_ICONS[step.tool_name] || "🔧"; - const args = step.arguments; - const dateRange = args.start_date && args.end_date ? `${args.start_date} → ${args.end_date}` : ""; - const filter = args.team_ids || args.tags || args.user_id || ""; - - return ( -
- - {step.status === "running" ? ( - - ) : step.status === "error" ? ( - ✗ - ) : ( - ✓ - )} - -
-
- {icon} {step.tool_label} -
- {dateRange &&
{dateRange}
} - {filter &&
Filter: {filter}
} - {step.status === "error" && step.error &&
{step.error}
} -
-
- ); -}; - -const MarkdownContent: React.FC<{ content: string }> = ({ content }) => ( -

{children}

, - strong: ({ children }) => {children}, - ul: ({ children }) =>
    {children}
, - ol: ({ children }) =>
    {children}
, - li: ({ children }) =>
  • {children}
  • , - h1: ({ children }) =>

    {children}

    , - h2: ({ children }) =>

    {children}

    , - h3: ({ children }) =>

    {children}

    , - code: ({ children, className }) => { - const isBlock = className?.includes("language-"); - return isBlock ? ( -
    -            {children}
    -          
    - ) : ( - {children} - ); - }, - table: ({ children }) => ( -
    - {children}
    -
    - ), - th: ({ children }) => ( - {children} - ), - td: ({ children }) => {children}, - }} - > - {content} -
    -); - -const UsageAIChatPanel: React.FC = ({ open, onClose, accessToken }) => { - const [messages, setMessages] = useState([]); - const [inputText, setInputText] = useState(""); - const [isLoading, setIsLoading] = useState(false); - const [selectedModel, setSelectedModel] = useState(undefined); - const [availableModels, setAvailableModels] = useState([]); - const [isLoadingModels, setIsLoadingModels] = useState(false); - const [streamingContent, setStreamingContent] = useState(""); - const [statusMessage, setStatusMessage] = useState(null); - const [activeToolCalls, setActiveToolCalls] = useState([]); - const messagesEndRef = useRef(null); - const abortControllerRef = useRef(null); - - useEffect(() => { - if (open && availableModels.length === 0) { - loadModels(); - } - }, [open]); - - useEffect(() => { - if (typeof messagesEndRef.current?.scrollIntoView === "function") { - messagesEndRef.current.scrollIntoView({ behavior: "smooth" }); - } - }, [messages, streamingContent, activeToolCalls, statusMessage]); - - const loadModels = async () => { - if (!accessToken) return; - setIsLoadingModels(true); - try { - const fetchedModels = await modelHubCall(accessToken); - if (fetchedModels?.data?.length > 0) { - const models = fetchedModels.data.map((item: any) => item.model_group as string).sort(); - setAvailableModels(models); - } - } catch (error) { - console.error("Failed to load models:", error); - } finally { - setIsLoadingModels(false); - } - }; - - const handleSend = async () => { - if (!accessToken || !inputText.trim() || isLoading) return; - - const userMessage: ChatMessage = { role: "user", content: inputText.trim() }; - const updatedMessages = [...messages, userMessage]; - setMessages(updatedMessages); - setInputText(""); - setIsLoading(true); - setStreamingContent(""); - setStatusMessage(null); - setActiveToolCalls([]); - - const abortController = new AbortController(); - abortControllerRef.current = abortController; - - let accumulated = ""; - const toolCalls: ToolCallStep[] = []; - - try { - await usageAiChatStream( - accessToken, - updatedMessages.slice(-20).map((m) => ({ role: m.role, content: m.content })), - selectedModel || "", - (content: string) => { - setStatusMessage(null); - accumulated += content; - setStreamingContent(accumulated); - }, - () => { - setStatusMessage(null); - setActiveToolCalls([]); - setMessages((prev) => [ - ...prev, - { role: "assistant", content: accumulated, toolCalls: toolCalls.length > 0 ? [...toolCalls] : undefined }, - ]); - setStreamingContent(""); - }, - (errorMsg: string) => { - setStatusMessage(null); - setActiveToolCalls([]); - setMessages((prev) => [...prev, { role: "assistant", content: `Error: ${errorMsg}` }]); - setStreamingContent(""); - }, - (status: string) => { - setStatusMessage(status); - }, - (event: UsageAiToolCallEvent) => { - const idx = toolCalls.findIndex((tc) => tc.tool_name === event.tool_name); - if (idx >= 0) { - toolCalls[idx] = { ...event }; - } else { - toolCalls.push({ ...event }); - } - setActiveToolCalls([...toolCalls]); - }, - abortController.signal, - ); - } catch (error: any) { - if (error?.name === "AbortError" || abortController.signal.aborted) { - return; - } - const errorMsg = error?.message || "Failed to get response. Please try again."; - setMessages((prev) => [...prev, { role: "assistant", content: `Error: ${errorMsg}` }]); - setStreamingContent(""); - } finally { - setIsLoading(false); - abortControllerRef.current = null; - } - }; - - const handleKeyDown = (e: React.KeyboardEvent) => { - if (e.key === "Enter" && !e.shiftKey) { - e.preventDefault(); - handleSend(); - } - }; - - const handleClose = () => { - if (abortControllerRef.current) { - abortControllerRef.current.abort(); - } - onClose(); - }; - - const handleClear = () => { - setMessages([]); - setStreamingContent(""); - setActiveToolCalls([]); - setStatusMessage(null); - }; - - return ( -
    - {/* Header */} -
    -
    -

    -

    -

    Ask about your spend, models, keys, and trends

    -
    - -
    - - {/* Model selector */} -
    - setSelectedModel(value ?? undefined)} - > - - - {isLoadingModels ? "Loading models…" : "No models found"} - - {(model: string) => ( - - {model} - - )} - - - -
    - - {/* Chat messages */} -
    - {messages.length === 0 && !streamingContent && !isLoading && ( -
    -
    - )} - - {messages.map((msg, idx) => ( -
    - {msg.role === "user" ? ( -
    -
    - {msg.content} -
    -
    - ) : ( -
    - {/* Tool calls for this message */} - {msg.toolCalls && msg.toolCalls.length > 0 && ( -
    - {msg.toolCalls.map((tc, tcIdx) => ( - - ))} -
    - )} - {/* Response */} -
    - -
    -
    - )} -
    - ))} - - {/* Active tool calls (in-progress) */} - {isLoading && activeToolCalls.length > 0 && ( -
    - {activeToolCalls.map((tc, idx) => ( - - ))} -
    - )} - - {/* Status / spinner */} - {isLoading && !streamingContent && ( -
    - - {statusMessage || "Thinking..."} -
    - )} - - {/* Streaming response */} - {streamingContent && ( -
    - -
    - )} - -
    -
    - - {/* Input area */} -
    -
    -