diff --git a/.circleci/config.yml b/.circleci/config.yml index e171759f1c4..6c7bbddb9f1 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -112,6 +112,24 @@ jobs: python -m mypy . cd .. no_output_timeout: 10m + + semgrep: + docker: + - image: cimg/python:3.12 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + steps: + - checkout + - setup_google_dns + - run: + name: Install Semgrep + command: pip install semgrep + - run: + name: Run Semgrep (custom rules only) + command: semgrep scan --config .semgrep/rules . --error + local_testing_part1: docker: - image: cimg/python:3.12 @@ -2277,6 +2295,7 @@ jobs: - run: python ./tests/code_coverage_tests/router_code_coverage.py - run: python ./tests/code_coverage_tests/test_chat_completion_imports.py - run: python ./tests/code_coverage_tests/info_log_check.py + - run: python ./tests/code_coverage_tests/check_guardrail_apply_decorator.py - run: python ./tests/code_coverage_tests/test_ban_set_verbose.py - run: python ./tests/code_coverage_tests/code_qa_check_tests.py - run: python ./tests/code_coverage_tests/check_get_model_cost_key_performance.py @@ -3801,7 +3820,6 @@ jobs: - run: name: Get new version command: | - cd litellm-proxy-extras NEW_VERSION=$(python -c "import toml; print(toml.load('pyproject.toml')['tool']['poetry']['version'])") echo "export NEW_VERSION=$NEW_VERSION" >> $BASH_ENV @@ -3826,7 +3844,6 @@ jobs: - run: name: Publish to PyPI command: | - cd litellm-proxy-extras echo -e "[pypi]\nusername = $PYPI_PUBLISH_USERNAME\npassword = $PYPI_PUBLISH_PASSWORD" > ~/.pypirc python -m pip install --upgrade pip build twine setuptools wheel rm -rf build dist @@ -3933,6 +3950,9 @@ jobs: image: ubuntu-2204:2023.10.1 resource_class: xlarge working_directory: ~/project + parameters: + browser: + type: string steps: - checkout - setup_google_dns @@ -3962,7 +3982,7 @@ jobs: echo "Expires at: $EXPIRES_AT" neon branches create \ --project-id $NEON_PROJECT_ID \ - --name preview/commit-${CIRCLE_SHA1:0:7} \ + --name preview/commit-${CIRCLE_SHA1:0:7}-<< parameters.browser >> \ --expires-at $EXPIRES_AT \ --parent br-fancy-paper-ad1olsb3 \ --api-key $NEON_API_KEY || true @@ -3972,7 +3992,7 @@ jobs: E2E_UI_TEST_DATABASE_URL=$(neon connection-string \ --project-id $NEON_PROJECT_ID \ --api-key $NEON_API_KEY \ - --branch preview/commit-${CIRCLE_SHA1:0:7} \ + --branch preview/commit-${CIRCLE_SHA1:0:7}-<< parameters.browser >> \ --database-name yuneng-trial-db \ --role neondb_owner) echo $E2E_UI_TEST_DATABASE_URL @@ -3984,7 +4004,7 @@ jobs: -e UI_USERNAME="admin" \ -e UI_PASSWORD="gm" \ -e LITELLM_LICENSE=$LITELLM_LICENSE \ - --name litellm-docker-database \ + --name litellm-docker-database-<< parameters.browser >> \ -v $(pwd)/litellm/proxy/example_config_yaml/simple_config.yaml:/app/config.yaml \ litellm-docker-database:ci \ --config /app/config.yaml \ @@ -4000,7 +4020,7 @@ jobs: sudo rm dockerize-linux-amd64-v0.6.1.tar.gz - run: name: Start outputting logs - command: docker logs -f litellm-docker-database + command: docker logs -f litellm-docker-database-<< parameters.browser >> background: true - run: name: Wait for app to be ready @@ -4009,6 +4029,7 @@ jobs: name: Run Playwright Tests command: | npx playwright test \ + --project << parameters.browser >> \ --config ui/litellm-dashboard/e2e_tests/playwright.config.ts \ --reporter=html \ --output=test-results @@ -4115,6 +4136,12 @@ workflows: only: - main - /litellm_.*/ + - semgrep: + filters: + branches: + only: + - main + - /litellm_.*/ - local_testing_part1: filters: branches: @@ -4214,6 +4241,20 @@ workflows: - main - /litellm_.*/ - e2e_ui_testing: + name: e2e_ui_testing_chromium + browser: chromium + context: e2e_ui_tests + requires: + - ui_build + - build_docker_database_image + filters: + branches: + only: + - main + - /litellm_.*/ + - e2e_ui_testing: + name: e2e_ui_testing_firefox + browser: firefox context: e2e_ui_tests requires: - ui_build @@ -4493,6 +4534,7 @@ workflows: - publish_to_pypi: requires: - mypy_linting + - semgrep - local_testing_part1 - local_testing_part2 - build_and_test @@ -4525,7 +4567,8 @@ workflows: - litellm_assistants_api_testing - auth_ui_unit_tests - db_migration_disable_update_check - - e2e_ui_testing + - e2e_ui_testing_chromium + - e2e_ui_testing_firefox - litellm_proxy_unit_testing_key_generation - litellm_proxy_unit_testing_part1 - litellm_proxy_unit_testing_part2 diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index b91b16c955c..f13039f4516 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -9,6 +9,7 @@ - [ ] I have Added testing in the [`tests/litellm/`](https://github.com/BerriAI/litellm/tree/main/tests/litellm) directory, **Adding at least 1 test is a hard requirement** - [see details](https://docs.litellm.ai/docs/extras/contributing_code) - [ ] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code) - [ ] My PR's scope is as isolated as possible, it only solves 1 specific problem +- [ ] I have requested a Greptile review by commenting `@greptileai` and received a **Confidence Score of at least 4/5** before requesting a maintainer review ## CI (LiteLLM team) diff --git a/.semgrep/rules/README.md b/.semgrep/rules/README.md new file mode 100644 index 00000000000..6cffcc32963 --- /dev/null +++ b/.semgrep/rules/README.md @@ -0,0 +1,52 @@ +# Custom Semgrep Rules + +All `.yml` files under `.semgrep/rules/` run in CI (CircleCI `semgrep` job). + +## Add a Rule + +* Add a `.yml` file under `.semgrep/rules///` + + +[Rule syntax →](https://semgrep.dev/docs/writing-rules/rule-syntax/) + +## Organizing Rules + +### Structure: language → domain + +``` +.semgrep/rules///.yml +``` + +Examples: + +- `python/security/unsafe-yaml-load.yml` +- `python/reliability/missing-timeout-http.yml` +- `python/performance/blocking-io-in-async.yml` + +### Rule metadata + +Match tags to the folder for consistent filtering: + +```yaml +metadata: + tags: [python, security] +``` + +### Severity expectations + +All rules must fail CI on findings. No warn-only rules. + +- Use `severity: ERROR` in rule metadata +- If a rule is noisy → refine until low false positives before adding + +## Run Locally + +```bash +semgrep scan --config .semgrep/rules . --error +``` + +With Semgrep registry: + +```bash +semgrep scan --config auto --config .semgrep/rules . +``` diff --git a/.semgrep/rules/python/reliability/unbounded-memory.yml b/.semgrep/rules/python/reliability/unbounded-memory.yml new file mode 100644 index 00000000000..f13c38471fb --- /dev/null +++ b/.semgrep/rules/python/reliability/unbounded-memory.yml @@ -0,0 +1,17 @@ +# Unbounded memory growth – data structures without a clear max limit +# Can lead to OOM under load. + +rules: + - id: unbounded-asyncio-queue + message: asyncio.Queue() with no maxsize can grow unbounded. Use asyncio.Queue(maxsize=N) for integrations (e.g. log queues). + severity: ERROR + languages: [python] + pattern-either: + - pattern: asyncio.Queue() + - pattern: asyncio.Queue(maxsize=0) + metadata: + category: reliability + cwe: "CWE-400: Uncontrolled Resource Consumption" + tags: [python, reliability] + confidence: HIGH + source: https://docs.python.org/3/library/asyncio-queue.html diff --git a/ci_cd/.grype.yaml b/ci_cd/.grype.yaml index 642e2dd9d03..b9bc9db58f5 100644 --- a/ci_cd/.grype.yaml +++ b/ci_cd/.grype.yaml @@ -1,3 +1,36 @@ ignore: - vulnerability: CVE-2026-22184 reason: no fixed zlib package is available yet in the Wolfi repositories, so this is ignored temporarily until an upstream release exists + # Wolfi base image: Python 3.13 and Node from apk have no fixed builds in Wolfi yet / not applicable + - vulnerability: CVE-2025-55130 + reason: Node in Wolfi apk; only used for Admin UI build/prisma + - vulnerability: CVE-2025-59465 + reason: Node in Wolfi apk; only used for Admin UI build/prisma + - vulnerability: CVE-2025-55131 + reason: Node in Wolfi apk; only used for Admin UI build/prisma + - vulnerability: CVE-2025-59466 + reason: Node in Wolfi apk; only used for Admin UI build/prisma + - vulnerability: CVE-2026-21637 + reason: Node in Wolfi apk; only used for Admin UI build/prisma + - vulnerability: CVE-2025-55132 + reason: Node in Wolfi apk; only used for Admin UI build/prisma + - vulnerability: GHSA-hx9q-6w63-j58v + reason: orjson dumps recursion; allowlisted + - vulnerability: GHSA-73rr-hh4g-fpgx + reason: diff npm transitive dep; override in package.json, allowlisted + - vulnerability: CVE-2026-0865 + reason: Python 3.13 in Wolfi base; no fixed apk build yet + - vulnerability: CVE-2025-15282 + reason: Python 3.13 in Wolfi base; no fixed apk build yet + - vulnerability: CVE-2026-0672 + reason: Python 3.13 in Wolfi base; no fixed apk build yet + - vulnerability: CVE-2025-15366 + reason: Python 3.13 in Wolfi base; no fixed apk build yet + - vulnerability: CVE-2025-15367 + reason: Python 3.13 in Wolfi base; no fixed apk build yet + - vulnerability: CVE-2025-11468 + reason: Python 3.13 in Wolfi base; no fixed apk build yet + - vulnerability: CVE-2025-12781 + reason: Python 3.13 in Wolfi base; no fixed apk build yet + - vulnerability: CVE-2026-1299 + reason: Python 3.13 in Wolfi base; no fixed apk build yet diff --git a/ci_cd/security_scans.sh b/ci_cd/security_scans.sh index 3ffa13c444f..2db72ae5c69 100755 --- a/ci_cd/security_scans.sh +++ b/ci_cd/security_scans.sh @@ -140,12 +140,14 @@ run_grype_scans() { "GHSA-34x7-hfp2-rc4v" # node-tar hardlink path traversal - not applicable, tar CLI not exposed in application code "GHSA-r6q2-hw4h-h46w" # node-tar not used by application runtime, Linux-only container, not affect by macOS APFS-specific exploit "GHSA-8rrh-rw8j-w5fx" # wheel is from chainguard and will be handled by then TODO: Remove this after Chainguard updates the wheel - "CVE-2025-59465" # We do not use Node in application runtime, only used for building Admin UI - "CVE-2025-55131" # We do not use Node in application runtime, only used for building Admin UI - "CVE-2025-59466" # We do not use Node in application runtime, only used for building Admin UI - "CVE-2025-55130" # We do not use Node in application runtime, only used for building Admin UI - "CVE-2025-59467" # We do not use Node in application runtime, only used for building Admin UI - "CVE-2026-21637" # We do not use Node in application runtime, only used for building Admin UI + "CVE-2025-59465" # Node only used for Admin UI build/prisma + "CVE-2025-55131" # Node only used for Admin UI build/prisma + "CVE-2025-59466" # Node only used for Admin UI build/prisma + "CVE-2025-55130" # Node only used for Admin UI build/prisma + "CVE-2025-59467" # Node only used for Admin UI build/prisma + "CVE-2026-21637" # Node only used for Admin UI build/prisma + "CVE-2025-55132" # Node only used for Admin UI build/prisma + "GHSA-hx9q-6w63-j58v" # orjson dumps recursion; allowlisted "CVE-2025-15281" # No fix available yet "CVE-2026-0865" # No fix available yet "CVE-2025-15282" # No fix available yet diff --git a/cookbook/nova_sonic_realtime.py b/cookbook/nova_sonic_realtime.py index 0ea0badfb01..c7a73c1d00f 100644 --- a/cookbook/nova_sonic_realtime.py +++ b/cookbook/nova_sonic_realtime.py @@ -16,10 +16,14 @@ Usage: import asyncio import base64 import json +import os import pyaudio import websockets from typing import Optional +# Bounded queue size for audio chunks (configurable via env to avoid unbounded memory) +AUDIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 10_000)) + # Audio configuration (matching Nova Sonic requirements) INPUT_SAMPLE_RATE = 16000 # Nova Sonic expects 16kHz input OUTPUT_SAMPLE_RATE = 24000 # Nova Sonic outputs 24kHz @@ -40,7 +44,7 @@ class RealtimeClient: self.api_key = api_key self.ws: Optional[websockets.WebSocketClientProtocol] = None self.is_active = False - self.audio_queue = asyncio.Queue() + self.audio_queue = asyncio.Queue(maxsize=AUDIO_QUEUE_MAXSIZE) self.pyaudio = pyaudio.PyAudio() self.input_stream = None self.output_stream = None diff --git a/docs/my-website/blog/model_cost_map_incident/index.md b/docs/my-website/blog/model_cost_map_incident/index.md new file mode 100644 index 00000000000..b9ff20e4128 --- /dev/null +++ b/docs/my-website/blog/model_cost_map_incident/index.md @@ -0,0 +1,95 @@ +--- +slug: model-cost-map-incident +title: "Incident Report: Invalid model cost map on main" +date: 2026-02-10T10:00:00 +authors: + - name: Ishaan Jaffer + title: "CTO, LiteLLM" + url: https://www.linkedin.com/in/ishaanjaffer/ + image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg +tags: [incident-report, stability] +hide_table_of_contents: false +--- + +**Date:** January 27, 2026 +**Duration:** ~20 minutes +**Severity:** Low +**Status:** Resolved + +## Summary + +A malformed JSON entry in `model_prices_and_context_window.json` was merged to `main` ([`562f0a0`](https://github.com/BerriAI/litellm/commit/562f0a028251750e3d75386bee0e630d9796d0df)). This caused LiteLLM to silently fall back to a stale local copy of the model cost map. Users on older package versions lost cost tracking for newer models only (e.g. `azure/gpt-5.2`). No LLM calls were blocked. + +- **LLM calls and proxy routing:** No impact. +- **Cost tracking:** Impacted for newer models not present in the local backup. Older models were unaffected. The incident lasted ~20 minutes until the commit was reverted. + +{/* truncate */} + +--- + +## Background + +The model cost map is not in the request path. It is used after the LLM response comes back, inside a try/catch, to calculate spend. A missing entry never blocks a call. + +```mermaid +flowchart TD + A["1. litellm.completion() receives request + litellm/main.py"] --> B["2. Route to provider + litellm/litellm_core_utils/get_llm_provider_logic.py"] + B --> C["3. LLM returns response + litellm/main.py"] + C --> D["4. Post-call: look up model in cost map + litellm/cost_calculator.py"] + D -->|"found"| E["5a. Attach cost to response"] + D -->|"not found (try/catch)"| F["5b. Log warning, set cost=0"] + E --> G["6. Return response to caller"] + F --> G + + style D fill:#fff3cd,stroke:#ffc107 + style F fill:#fff3cd,stroke:#ffc107 + style E fill:#d4edda,stroke:#28a745 + style G fill:#d4edda,stroke:#28a745 +``` + +Both paths return a response to the caller. When the cost map lookup fails, the only difference is `cost=0` on that request. + +--- + +## Root cause + +LiteLLM fetches the model cost map from GitHub `main` at import time. If the fetch fails, it falls back to a local backup bundled with the package. Before this incident, the fallback was completely silent -- no warning was logged. + +A contributor PR introduced an extra `{` bracket, producing invalid JSON. The remote fetch failed with `JSONDecodeError`, triggering the silent fallback. Users on older package versions had backup files missing newer models. + +**Timeline:** + +1. Malformed JSON merged to `main` +2. LiteLLM installations fall back to local backup on next import +3. Users report `"This model isn't mapped yet"` for newer models +4. Bad commit identified and reverted (~20 minutes) + +--- + +## Remediation + +| # | Action | Status | Code | +|---|---|---|---| +| 1 | CI validation on `model_prices_and_context_window.json` | ✅ Done | [`test-model-map.yaml`](https://github.com/BerriAI/litellm/blob/main/.github/workflows/test-model-map.yaml) | +| 2 | Warning log on fallback to local backup | ✅ Done | [`get_model_cost_map.py#L57-L68`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm_core_utils/get_model_cost_map.py#L57-L68) | +| 3 | `GetModelCostMap` class with integrity validation helpers | ✅ Done | [`get_model_cost_map.py#L24-L149`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm_core_utils/get_model_cost_map.py#L24-L149) | +| 4 | Resilience test suite (bad hosted map, fallback, completion) | ✅ Done | [`test_model_cost_map_resilience.py#L150-L291`](https://github.com/BerriAI/litellm/blob/main/tests/llm_translation/test_model_cost_map_resilience.py#L150-L291) | +| 5 | Test that backup model cost map always exists and contains common models | ✅ Done | [`test_model_cost_map_resilience.py#L213-L228`](https://github.com/BerriAI/litellm/blob/main/tests/llm_translation/test_model_cost_map_resilience.py#L213-L228) | + +Enterprises that require zero external dependencies at import time can set `LITELLM_LOCAL_MODEL_COST_MAP=True` to skip the GitHub fetch entirely. + +--- + +## Other dependencies on external resources + +| Dependency | Impact if unavailable | Fallback | +|---|---|---| +| Model cost map (GitHub) | Cost tracking for newer models | Local backup (now with warning) | +| JWT public keys (IDP/SSO) | Auth fails | None | +| OIDC UserInfo (IDP/SSO) | Auth fails | None | +| HuggingFace model API | HF provider calls fail | None | +| Ollama tags (localhost) | Ollama model list stale | Static list | diff --git a/docs/my-website/docs/benchmarks.md b/docs/my-website/docs/benchmarks.md index a1489081b4c..1f818cef498 100644 --- a/docs/my-website/docs/benchmarks.md +++ b/docs/my-website/docs/benchmarks.md @@ -5,6 +5,13 @@ import Image from '@theme/IdealImage'; Benchmarks for LiteLLM Gateway (Proxy Server) tested against a fake OpenAI endpoint. +## Setting Up a Fake OpenAI Endpoint + +For load testing and benchmarking, you can use a fake OpenAI proxy server. LiteLLM provides: + +1. **Hosted endpoint**: Use our free hosted fake endpoint at `https://exampleopenaiendpoint-production.up.railway.app/` +2. **Self-hosted**: Set up your own fake OpenAI proxy server using [github.com/BerriAI/example_openai_endpoint](https://github.com/BerriAI/example_openai_endpoint) + Use this config for testing: ```yaml @@ -12,7 +19,7 @@ model_list: - model_name: "fake-openai-endpoint" litellm_params: model: openai/any - api_base: https://your-fake-openai-endpoint.com/chat/completions + api_base: https://exampleopenaiendpoint-production.up.railway.app/ # or your self-hosted endpoint api_key: "test" ``` diff --git a/docs/my-website/docs/load_test.md b/docs/my-website/docs/load_test.md index 4641a70366c..071b097904b 100644 --- a/docs/my-website/docs/load_test.md +++ b/docs/my-website/docs/load_test.md @@ -4,8 +4,9 @@ import Image from '@theme/IdealImage'; ## Locust Load Test LiteLLM Proxy -1. Add `fake-openai-endpoint` to your proxy config.yaml and start your litellm proxy -litellm provides a free hosted `fake-openai-endpoint` you can load test against +1. Add `fake-openai-endpoint` to your proxy config.yaml and start your litellm proxy. + +LiteLLM provides a free hosted `fake-openai-endpoint` you can load test against. You can also self-host your own fake OpenAI proxy server using [github.com/BerriAI/example_openai_endpoint](https://github.com/BerriAI/example_openai_endpoint). ```yaml model_list: diff --git a/docs/my-website/docs/load_test_advanced.md b/docs/my-website/docs/load_test_advanced.md index 3171bc33594..d35b5f74784 100644 --- a/docs/my-website/docs/load_test_advanced.md +++ b/docs/my-website/docs/load_test_advanced.md @@ -29,12 +29,16 @@ Tutorial on how to get to 1K+ RPS with LiteLLM Proxy on locust **Note:** we're currently migrating to aiohttp which has 10x higher throughput. We recommend using the `openai/` provider for load testing. +:::tip Setting Up a Fake OpenAI Endpoint +You can use our hosted fake endpoint or self-host your own using [github.com/BerriAI/example_openai_endpoint](https://github.com/BerriAI/example_openai_endpoint). +::: + ```yaml model_list: - model_name: "fake-openai-endpoint" litellm_params: model: openai/any - api_base: https://your-fake-openai-endpoint.com/chat/completions + api_base: https://exampleopenaiendpoint-production.up.railway.app/ # or your self-hosted endpoint api_key: "test" ``` diff --git a/docs/my-website/docs/mcp_oauth.md b/docs/my-website/docs/mcp_oauth.md index ed69408196f..9cd7b1e77be 100644 --- a/docs/my-website/docs/mcp_oauth.md +++ b/docs/my-website/docs/mcp_oauth.md @@ -1,6 +1,3 @@ -import Tabs from '@theme/Tabs'; -import TabItem from '@theme/TabItem'; - # MCP OAuth LiteLLM supports two OAuth 2.0 flows for MCP servers: @@ -98,8 +95,71 @@ LiteLLM automatically fetches, caches, and refreshes OAuth2 tokens using the `cl ### Setup - - +You can configure M2M OAuth via the LiteLLM UI or `config.yaml`. + +### UI Setup + +Navigate to the **MCP Servers** page and click **+ Add New MCP Server**. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/d1f1e89c-a789-4975-8846-b15d9821984a/ascreenshot_630800e00a2e4b598baabfc25efbabd3_text_export.jpeg) + +Enter a name for your server and select **HTTP** as the transport type. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/2008c9d6-6093-4121-beab-1e52c71376aa/ascreenshot_516ffd6c7b524465a253a56048c3d228_text_export.jpeg) + +Paste the MCP server URL. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/b0ee8b7d-6de8-492b-8962-287987feec29/ascreenshot_b3efca82078a4c6bb1453c58161909f9_text_export.jpeg) + +Under **Authentication**, select **OAuth**. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/e1597814-ff8e-40b9-9d7b-864dcdbe0910/ascreenshot_2097612712264d8f9e553f7ca9175fb0_text_export.jpeg) + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/f6ea5694-f28a-4bc3-9c9a-bb79f199bd65/ascreenshot_9be839f55b1b4f96bfe24030ba2c7f8d_text_export.jpeg) + +Choose **Machine-to-Machine (M2M)** as the OAuth flow type. This is for server-to-server authentication using the `client_credentials` grant — no browser interaction required. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/9853310c-1d86-4628-bad1-7a391eca0e4d/ascreenshot_f302a286fa264fdd8d56db53b8f9395c_text_export.jpeg) + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/df64dc65-ef86-475d-adaf-12e227d5e873/ascreenshot_9e2f41d43a76435f918a00b52ffcc639_text_export.jpeg) + +Fill in the **Client ID** and **Client Secret** provided by your OAuth provider. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/0de5a7bd-9898-4fc7-8843-b23dd5aac47f/ascreenshot_b9087aaa81a14b5b9c199929efc4a563_text_export.jpeg) + +Enter the **Token URL** — this is the endpoint LiteLLM will call to fetch access tokens using `client_credentials`. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/0aea70f1-558c-4dca-91bc-1175fe1ddc89/ascreenshot_b3fcf8a1287e4e2d9a3d67c4a29f7bff_text_export.jpeg) + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/e842ef09-1fd7-47a6-909b-252d389f0abc/ascreenshot_2a87dad3624847e7ac370591d1d1aedd_text_export.jpeg) + +Scroll down and review the server URL and all fields, then click **Create MCP Server**. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/0857712b-4b53-40f8-8c1f-a4c72edaa644/ascreenshot_47be3fcd5de64ed391f70c1fb74a8bfc_text_export.jpeg) + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/9d961765-955f-4905-a3dc-1a446aa3b2cc/ascreenshot_43fd39d014224564bc6b35aced1fb6d3_text_export.jpeg) + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/3825d5fa-8fd1-4e71-b090-77ff0259c3f6/ascreenshot_2509a7ebd9bf421eb0e82f2553566745_text_export.jpeg) + +Once created, open the server and navigate to the **MCP Tools** tab to verify that LiteLLM can connect and list available tools. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/8107e27b-5072-4675-8fd6-89b47692b1bd/ascreenshot_f774bc76138f430d808fb4482ebfcdca_text_export.jpeg) + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/ce94bb7b-c81b-4396-9939-178efb2cdfce/ascreenshot_28b838ab6ae34c76858454555c4c1d79_text_export.jpeg) + +Select a tool (e.g. **echo**) to test it. Fill in the required parameters and click **Call Tool**. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/c459c1d3-ec29-4211-9c28-37fbe7783bbc/ascreenshot_e9b138b3c2cc4440bb1a6f42ac7ae861_text_export.jpeg) + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/5438ac60-e0ac-4a79-bf6f-5594f160d3b5/ascreenshot_9133a17d26204c46bce497e74685c483_text_export.jpeg) + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/a8f6821b-3982-4b4d-9b25-70c8aff5ac31/ascreenshot_28d474d0e62545a482cff6128527883a_text_export.jpeg) + +LiteLLM automatically fetches an OAuth token behind the scenes and calls the tool. The result confirms the M2M OAuth flow is working end-to-end. + +![](https://colony-recorder.s3.amazonaws.com/files/2026-02-10/c3924549-a949-48d1-ac67-ab4c30475859/ascreenshot_8f6eca9d717f45478d50a881bd244bb3_text_export.jpeg) + +### Config.yaml Setup ```yaml title="config.yaml" showLineNumbers mcp_servers: @@ -112,14 +172,6 @@ mcp_servers: scopes: ["mcp:read", "mcp:write"] # optional ``` - - - -Navigate to **MCP Servers → Add Server → Authentication → OAuth**, then fill in `client_id`, `client_secret`, and `token_url`. - - - - ### How It Works 1. On first MCP request, LiteLLM POSTs to `token_url` with `grant_type=client_credentials` diff --git a/docs/my-website/docs/pass_through/bedrock.md b/docs/my-website/docs/pass_through/bedrock.md index b8d20d77da0..65c5d8caadc 100644 --- a/docs/my-website/docs/pass_through/bedrock.md +++ b/docs/my-website/docs/pass_through/bedrock.md @@ -556,3 +556,147 @@ for event in response.get("completion"): print(completion) ``` + +## Using LangChain AWS SDK with LiteLLM + +You can use the [LangChain AWS SDK](https://python.langchain.com/docs/integrations/chat/bedrock/) with LiteLLM Proxy to get cost tracking, load balancing, and other LiteLLM features. + +### Quick Start + +**1. Install LangChain AWS**: + +```bash showLineNumbers +pip install langchain-aws +``` + +**2. Setup LiteLLM Proxy**: + +Create a `config.yaml`: + +```yaml showLineNumbers +model_list: + - model_name: claude-sonnet + litellm_params: + model: bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0 + aws_region_name: us-east-1 + custom_llm_provider: bedrock +``` + +Start the proxy: + +```bash showLineNumbers +export AWS_ACCESS_KEY_ID="your-access-key" +export AWS_SECRET_ACCESS_KEY="your-secret-key" + +litellm --config config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +**3. Use LangChain with LiteLLM**: + +```python showLineNumbers +from langchain_aws import ChatBedrockConverse +from langchain_core.messages import HumanMessage + +# Your LiteLLM API key +API_KEY = "Bearer sk-1234" + +# Initialize ChatBedrockConverse pointing to LiteLLM proxy +llm = ChatBedrockConverse( + model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + endpoint_url="http://localhost:4000/bedrock", + region_name="us-east-1", + aws_access_key_id=API_KEY, + aws_secret_access_key="bedrock" # Any non-empty value works +) + +# Invoke the model +messages = [HumanMessage(content="Hello, how are you?")] +response = llm.invoke(messages) + +print(response.content) +``` + +### Advanced Example: PDF Document Processing with Citations + +LangChain AWS SDK supports Bedrock's document processing features. Here's how to use it with LiteLLM: + +```python showLineNumbers +import os +import json +from langchain_aws import ChatBedrockConverse +from langchain_core.messages import HumanMessage + +# Your LiteLLM API key +API_KEY = "Bearer sk-1234" + +def get_llm() -> ChatBedrockConverse: + """Initialize LLM pointing to LiteLLM proxy""" + llm = ChatBedrockConverse( + model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + base_model_id="anthropic.claude-3-7-sonnet-20250219-v1:0", + endpoint_url="http://localhost:4000/bedrock", + region_name="us-east-1", + aws_access_key_id=API_KEY, + aws_secret_access_key="bedrock" + ) + return llm + +if __name__ == "__main__": + # Initialize the LLM + llm = get_llm() + + # Read PDF file as bytes (Converse API requires raw bytes) + with open("your-document.pdf", "rb") as file: + file_bytes = file.read() + + # Prepare messages with document attachment + messages = [ + HumanMessage(content=[ + {"text": "What is the policy number in this document?"}, + { + "document": { + "format": "pdf", + "name": "PolicyDocument", + "source": {"bytes": file_bytes}, + "citations": {"enabled": True} + } + } + ]) + ] + + # Invoke the LLM + response = llm.invoke(messages) + + # Print response with citations + print(json.dumps(response.content, indent=4)) +``` + +### Supported LangChain Features + +All LangChain AWS features work with LiteLLM: + +| Feature | Supported | Notes | +|---------|-----------|-------| +| Text Generation | ✅ | Full support | +| Streaming | ✅ | Use `stream()` method | +| Document Processing | ✅ | PDF, images, etc. | +| Citations | ✅ | Enable in document config | +| Tool Use | ✅ | Function calling support | +| Multi-modal | ✅ | Text + images + documents | + +### Troubleshooting + +**Issue**: `UnknownOperationException` error + +**Solution**: Make sure you're using the correct endpoint URL format: +- ✅ Correct: `http://localhost:4000/bedrock` +- ❌ Wrong: `http://localhost:4000/bedrock/v2` + +**Issue**: Authentication errors + +**Solution**: Ensure your API key is in the correct format: +```python +aws_access_key_id="Bearer sk-1234" # Include "Bearer " prefix +``` diff --git a/docs/my-website/docs/providers/perplexity.md b/docs/my-website/docs/providers/perplexity.md index 2fcb49c60fa..68adf9939c6 100644 --- a/docs/my-website/docs/providers/perplexity.md +++ b/docs/my-website/docs/providers/perplexity.md @@ -120,6 +120,293 @@ All models listed here https://docs.perplexity.ai/docs/model-cards are supported +## Agentic Research API (Responses API) + +Requires v1.72.6+ + + +### Using Presets + +Presets provide optimized defaults for specific use cases. Start with a preset for quick setup: + + + + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +# Using the pro-search preset +response = responses( + model="perplexity/preset/pro-search", + input="What are the latest developments in AI?", + custom_llm_provider="perplexity", +) + +print(response.output) +``` + + + + +1. Setup config.yaml + +```yaml +model_list: + - model_name: perplexity-pro-search + litellm_params: + model: perplexity/preset/pro-search + api_key: os.environ/PERPLEXITY_API_KEY +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Test it! + +```bash +curl http://0.0.0.0:4000/v1/responses \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer anything" \ + -d '{ + "model": "perplexity-pro-search", + "input": "What are the latest developments in AI?" + }' +``` + + + + +### Using Third-Party Models + +Access models from OpenAI, Anthropic, Google, xAI, and other providers through Perplexity's unified API: + + + + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/openai/gpt-4o", + input="Explain quantum computing in simple terms", + custom_llm_provider="perplexity", + max_output_tokens=500, +) + +print(response.output) +``` + + + + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/anthropic/claude-3-5-sonnet-20241022", + input="Write a short story about a robot learning to paint", + custom_llm_provider="perplexity", + max_output_tokens=500, +) + +print(response.output) +``` + + + + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/google/gemini-2.0-flash-exp", + input="Explain the concept of neural networks", + custom_llm_provider="perplexity", + max_output_tokens=500, +) + +print(response.output) +``` + + + + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/xai/grok-2-1212", + input="What makes a good AI assistant?", + custom_llm_provider="perplexity", + max_output_tokens=500, +) + +print(response.output) +``` + + + + +### Web Search Tool + +Enable web search capabilities to access real-time information: + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/openai/gpt-4o", + input="What's the weather in San Francisco today?", + custom_llm_provider="perplexity", + tools=[{"type": "web_search"}], + instructions="You have access to a web_search tool. Use it for questions about current events.", +) + +print(response.output) +``` + + +### Reasoning Effort (Responses API) + +Control the reasoning effort level for reasoning-capable models: + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/openai/gpt-5.2", + input="Solve this complex problem step by step", + custom_llm_provider="perplexity", + reasoning={"effort": "high"}, # Options: low, medium, high + max_output_tokens=1000, +) + +print(response.output) +``` + +### Multi-Turn Conversations + +Use message arrays for multi-turn conversations with context: + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/anthropic/claude-3-5-sonnet-20241022", + input=[ + {"type": "message", "role": "system", "content": "You are a helpful assistant."}, + {"type": "message", "role": "user", "content": "What are the latest AI developments?"}, + ], + custom_llm_provider="perplexity", + instructions="Provide detailed, well-researched answers.", + max_output_tokens=800, +) + +print(response.output) +``` + +### Streaming Responses + +Stream responses for real-time output: + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +response = responses( + model="perplexity/openai/gpt-4o", + input="Tell me a story about space exploration", + custom_llm_provider="perplexity", + stream=True, + max_output_tokens=500, +) + +for chunk in response: + if hasattr(chunk, 'type'): + if chunk.type == "response.output_text.delta": + print(chunk.delta, end="", flush=True) +``` + +### Supported Third-Party Models + +| Provider | Model Name | Function Call | +|----------|------------|---------------| +| OpenAI | gpt-4o | `responses(model="perplexity/openai/gpt-4o", ...)` | +| OpenAI | gpt-4o-mini | `responses(model="perplexity/openai/gpt-4o-mini", ...)` | +| OpenAI | gpt-5.2 | `responses(model="perplexity/openai/gpt-5.2", ...)` | +| Anthropic | claude-3-5-sonnet-20241022 | `responses(model="perplexity/anthropic/claude-3-5-sonnet-20241022", ...)` | +| Anthropic | claude-3-5-haiku-20241022 | `responses(model="perplexity/anthropic/claude-3-5-haiku-20241022", ...)` | +| Google | gemini-2.0-flash-exp | `responses(model="perplexity/google/gemini-2.0-flash-exp", ...)` | +| Google | gemini-2.0-flash-thinking-exp | `responses(model="perplexity/google/gemini-2.0-flash-thinking-exp", ...)` | +| xAI | grok-2-1212 | `responses(model="perplexity/xai/grok-2-1212", ...)` | +| xAI | grok-2-vision-1212 | `responses(model="perplexity/xai/grok-2-vision-1212", ...)` | + +### Available Presets + +| Preset Name | Function Call | +|----------------|--------------------------------------------------------| +| fast-search | `responses(model="perplexity/preset/fast-search", ...)`| +| pro-search | `responses(model="perplexity/preset/pro-search", ...)` | +| deep-research | `responses(model="perplexity/preset/deep-research", ...)`| + +### Complete Example + +```python +from litellm import responses +import os + +os.environ['PERPLEXITY_API_KEY'] = "" + +# Comprehensive example with multiple features +response = responses( + model="perplexity/openai/gpt-4o", + input="Research the latest developments in quantum computing and provide sources", + custom_llm_provider="perplexity", + tools=[ + {"type": "web_search"}, + {"type": "fetch_url"} + ], + instructions="Use web_search to find relevant information and fetch_url to retrieve detailed content from sources. Provide citations for all claims.", + max_output_tokens=1000, + temperature=0.7, +) + +print(f"Response ID: {response.id}") +print(f"Model: {response.model}") +print(f"Status: {response.status}") +print(f"Output: {response.output}") +print(f"Usage: {response.usage}") +``` + :::info For more information about passing provider-specific parameters, [go here](../completion/provider_specific_params.md) diff --git a/docs/my-website/docs/proxy/admin_ui_sso.md b/docs/my-website/docs/proxy/admin_ui_sso.md index 37e45b50284..f88d3480446 100644 --- a/docs/my-website/docs/proxy/admin_ui_sso.md +++ b/docs/my-website/docs/proxy/admin_ui_sso.md @@ -223,6 +223,7 @@ GENERIC_USER_FIRST_NAME_ATTRIBUTE = "first_name" GENERIC_USER_LAST_NAME_ATTRIBUTE = "last_name" GENERIC_USER_ROLE_ATTRIBUTE = "given_role" GENERIC_USER_PROVIDER_ATTRIBUTE = "provider" +GENERIC_USER_EXTRA_ATTRIBUTES = "department,employee_id,manager" # comma-separated list of additional fields to extract from SSO response GENERIC_CLIENT_STATE = "some-state" # if the provider needs a state parameter GENERIC_INCLUDE_CLIENT_ID = "false" # some providers enforce that the client_id is not in the body GENERIC_SCOPE = "openid profile email" # default scope openid is sometimes not enough to retrieve basic user info like first_name and last_name located in profile scope @@ -239,6 +240,40 @@ Use `GENERIC_USER_ROLE_ATTRIBUTE` to specify which attribute in the SSO token co Nested attribute paths are supported (e.g., `claims.role` or `attributes.litellm_role`). +**Capturing Additional SSO Fields** + +Use `GENERIC_USER_EXTRA_ATTRIBUTES` to extract additional fields from the SSO provider response beyond the standard user attributes (id, email, name, etc.). This is useful when you need to access custom organization-specific data (e.g., department, employee ID, groups) in your [custom SSO handler](./custom_sso.md). + +```shell +# Comma-separated list of field names to extract +GENERIC_USER_EXTRA_ATTRIBUTES="department,employee_id,manager,groups" +``` + +**Accessing Extra Fields in Custom SSO Handler:** + +```python +from litellm.proxy.management_endpoints.types import CustomOpenID + +async def custom_sso_handler(userIDPInfo: CustomOpenID): + # Access the extra fields + extra_fields = getattr(userIDPInfo, 'extra_fields', None) or {} + + user_department = extra_fields.get("department") + employee_id = extra_fields.get("employee_id") + user_groups = extra_fields.get("groups", []) + + # Use these fields for custom logic (e.g., team assignment, access control) + # ... +``` + +**Nested Field Paths:** + +Dot notation is supported for nested fields: + +```shell +GENERIC_USER_EXTRA_ATTRIBUTES="org_info.department,org_info.cost_center,metadata.employee_type" +``` + - Set Redirect URI, if your provider requires it - Set a redirect url = `/sso/callback` ```shell diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 5cdae51f448..38ad9bdd0ee 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -395,7 +395,7 @@ router_settings: | ATHINA_API_KEY | API key for Athina service | ATHINA_BASE_URL | Base URL for Athina service (defaults to `https://log.athina.ai`) | AUTH_STRATEGY | Strategy used for authentication (e.g., OAuth, API key) -| AUTO_REDIRECT_UI_LOGIN_TO_SSO | Flag to enable automatic redirect of UI login page to SSO when SSO is configured. Default is **true** +| AUTO_REDIRECT_UI_LOGIN_TO_SSO | Flag to enable automatic redirect of UI login page to SSO when SSO is configured. Default is **false** | AUDIO_SPEECH_CHUNK_SIZE | Chunk size for audio speech processing. Default is 1024 | ANTHROPIC_API_KEY | API key for Anthropic service | ANTHROPIC_API_BASE | Base URL for Anthropic API. Default is https://api.anthropic.com @@ -548,6 +548,10 @@ router_settings: | DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL | Default embedding model for MCP semantic tool filtering. Default is "text-embedding-3-small" | DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD | Default similarity threshold for MCP semantic tool filtering. Default is 0.3 | DEFAULT_MCP_SEMANTIC_FILTER_TOP_K | Default number of top results to return for MCP semantic tool filtering. Default is 10 +| MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL | Default TTL in seconds for MCP OAuth2 token cache. Default is 3600 +| MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE | Maximum number of entries in MCP OAuth2 token cache. Default is 200 +| MCP_OAUTH2_TOKEN_CACHE_MIN_TTL | Minimum TTL in seconds for MCP OAuth2 token cache. Default is 10 +| MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from token expiry when computing cache TTL. Default is 60 | DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT | Default token count for mock response completions. Default is 20 | DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT | Default token count for mock response prompts. Default is 10 | DEFAULT_MODEL_CREATED_AT_TIME | Default creation timestamp for models. Default is 1677610602 @@ -640,6 +644,7 @@ router_settings: | GENERIC_TOKEN_ENDPOINT | Token endpoint for generic OAuth providers | GENERIC_USER_DISPLAY_NAME_ATTRIBUTE | Attribute for user's display name in generic auth | GENERIC_USER_EMAIL_ATTRIBUTE | Attribute for user's email in generic auth +| GENERIC_USER_EXTRA_ATTRIBUTES | Comma-separated list of additional fields to extract from generic SSO provider response (e.g., "department,employee_id,groups"). Accessible via `CustomOpenID.extra_fields` in custom SSO handlers. Supports dot notation for nested fields | GENERIC_USER_FIRST_NAME_ATTRIBUTE | Attribute for user's first name in generic auth | GENERIC_USER_ID_ATTRIBUTE | Attribute for user ID in generic auth | GENERIC_USER_LAST_NAME_ATTRIBUTE | Attribute for user's last name in generic auth @@ -779,6 +784,7 @@ router_settings: | LITELLM_USER_AGENT | Custom user agent string for LiteLLM API requests. Used for partner telemetry attribution | LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD | If true, prints the standard logging payload to the console - useful for debugging | LITELM_ENVIRONMENT | Environment for LiteLLM Instance. This is currently only logged to DeepEval to determine the environment for DeepEval integration. +| LITELLM_ASYNCIO_QUEUE_MAXSIZE | Maximum size for asyncio queues (e.g. log queues, spend update queues, and cookbook examples such as realtime audio in `nova_sonic_realtime.py`). Bounds in-memory growth to prevent OOM. Default is 1000. | LOGFIRE_TOKEN | Token for Logfire logging service | LOGFIRE_BASE_URL | Base URL for Logfire logging service (useful for self hosted deployments) | LOGGING_WORKER_CONCURRENCY | Maximum number of concurrent coroutine slots for the logging worker on the asyncio event loop. Default is 100. Setting too high will flood the event loop with logging tasks which will lower the overall latency of the requests. @@ -806,6 +812,7 @@ router_settings: | MAX_RETRY_DELAY | Maximum delay in seconds for retrying requests. Default is 8.0 | MAX_LANGFUSE_INITIALIZED_CLIENTS | Maximum number of Langfuse clients to initialize on proxy. Default is 50. This is set since langfuse initializes 1 thread everytime a client is initialized. We've had an incident in the past where we reached 100% cpu utilization because Langfuse was initialized several times. | MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH | Maximum header length for MCP semantic filter tools. Default is 150 +| MAX_POLICY_ESTIMATE_IMPACT_ROWS | Maximum number of rows returned when estimating the impact of a policy. Default is 1000 | MIN_NON_ZERO_TEMPERATURE | Minimum non-zero temperature value. Default is 0.0001 | MINIMUM_PROMPT_CACHE_TOKEN_COUNT | Minimum token count for caching a prompt. Default is 1024 | MISTRAL_API_BASE | Base URL for Mistral API. Default is https://api.mistral.ai @@ -822,6 +829,8 @@ router_settings: | MICROSOFT_USER_ID_ATTRIBUTE | Field name for user ID in Microsoft SSO response. Default is `id` | MICROSOFT_USER_LAST_NAME_ATTRIBUTE | Field name for user last name in Microsoft SSO response. Default is `surname` | MICROSOFT_USERINFO_ENDPOINT | Custom userinfo endpoint URL for Microsoft SSO (overrides default Microsoft Graph userinfo endpoint) +| MODEL_COST_MAP_MAX_SHRINK_RATIO | Maximum allowed shrinkage ratio when validating a fetched model cost map against the local backup. Rejects the fetched map if it is smaller than this fraction of the backup. Default is 0.5 +| MODEL_COST_MAP_MIN_MODEL_COUNT | Minimum number of models a fetched cost map must contain to be considered valid. Default is 50 | NO_DOCS | Flag to disable Swagger UI documentation | NO_REDOC | Flag to disable Redoc documentation | NO_PROXY | List of addresses to bypass proxy diff --git a/docs/my-website/docs/proxy/custom_sso.md b/docs/my-website/docs/proxy/custom_sso.md index bbd7f41bee1..8b7adeb0c5a 100644 --- a/docs/my-website/docs/proxy/custom_sso.md +++ b/docs/my-website/docs/proxy/custom_sso.md @@ -142,6 +142,18 @@ async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: f"No ID found for user. userIDPInfo.id is None {userIDPInfo}" ) + ################################################# + # Access extra fields from SSO provider (requires GENERIC_USER_EXTRA_ATTRIBUTES env var) + # Example: Set GENERIC_USER_EXTRA_ATTRIBUTES="department,employee_id,groups" + extra_fields = getattr(userIDPInfo, 'extra_fields', None) or {} + user_department = extra_fields.get("department") + employee_id = extra_fields.get("employee_id") + user_groups = extra_fields.get("groups", []) + + print(f"User department: {user_department}") # noqa + print(f"Employee ID: {employee_id}") # noqa + print(f"User groups: {user_groups}") # noqa + ################################################# ################################################# # Run your custom code / logic here diff --git a/docs/my-website/docs/proxy/guardrails/guardrail_policies.md b/docs/my-website/docs/proxy/guardrails/guardrail_policies.md index 56be11c85a7..e2cb839203e 100644 --- a/docs/my-website/docs/proxy/guardrails/guardrail_policies.md +++ b/docs/my-website/docs/proxy/guardrails/guardrail_policies.md @@ -1,3 +1,7 @@ +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + # [Beta] Guardrail Policies Use policies to group guardrails and control which ones run for specific teams, keys, or models. @@ -10,6 +14,9 @@ Use policies to group guardrails and control which ones run for specific teams, ## Quick Start + + + ```yaml showLineNumbers title="config.yaml" model_list: - model_name: gpt-4 @@ -43,6 +50,26 @@ policy_attachments: scope: "*" # apply to all requests ``` + + + +**Step 1: Create a Policy** + +Go to **Policies** tab and click **+ Create New Policy**. Fill in the policy name, description, and select guardrails to add. + +![Enter policy name](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/4ba62cc8-d2c4-4af1-a526-686295466928/ascreenshot_401eab3e2081466e8f4d4ffa3bf7bff4_text_export.jpeg) + +![Add a description for the policy](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/51685e47-1d94-4d9c-acb0-3c88dce9f938/ascreenshot_a5cd40066ff34afbb1e4089a3c93d889_text_export.jpeg) + +![Select a parent policy to inherit from](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/1d96c3d3-187a-4f7c-97d2-6ac1f093d51e/ascreenshot_8a3af3b2210547dca3d4709df920d005_text_export.jpeg) + +![Select guardrails to add to the policy](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/23781274-e600-4d5f-a8a6-4a2a977a166c/ascreenshot_a2a45d2c5d064c77ab7cb47b569ad9e9_text_export.jpeg) + +![Click Create Policy to save](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/1d1ae8a8-daa5-451b-9fa2-c5b607ff6220/ascreenshot_218c2dd259714be4aa3c4e1894c96878_text_export.jpeg) + + + + Response headers show what ran: ``` @@ -58,6 +85,9 @@ x-litellm-applied-guardrails: pii_masking,prompt_injection You have a global baseline, but want to add extra guardrails for a specific team. + + + ```yaml showLineNumbers title="config.yaml" policies: global-baseline: @@ -81,6 +111,30 @@ policy_attachments: - finance # team alias from /team/new ``` + + + +**Option 1: Create a team-scoped attachment** + +Go to **Policies** > **Attachments** tab and click **+ Create New Attachment**. Select the policy and the teams to scope it to. + +![Select teams for the attachment](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/50e58f54-3bc3-477e-a106-e58cb65fde7e/ascreenshot_85d2e3d9d8d24842baced92fea170427_text_export.jpeg) + +![Select the teams to attach the policy to](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/f24066bb-0a73-49fb-87b6-c65ad3ca5b2f/ascreenshot_242476fbdac447309f65de78b0ed9fdd_text_export.jpeg) + +**Option 2: Attach from team settings** + +Go to **Teams** > click on a team > **Settings** tab > under **Policies**, select the policies to attach. + +![Open team settings and click Edit Settings](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/c31c3735-4f9d-4c6a-896b-186e97296940/ascreenshot_4749bb24ce5942cca462acc958fd3822_text_export.jpeg) + +![Select policies to attach to this team](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/da8d5d7a-d975-4bfe-acd2-f41dcea29520/ascreenshot_835a33b6cec545cbb2987f017fbaff90_text_export.jpeg) + + + + + + Now the `finance` team gets `pii_masking` + `strict_compliance_check` + `audit_logger`, while everyone else just gets `pii_masking`. ## Remove guardrails for a specific team @@ -201,6 +255,60 @@ policy_attachments: - "test-*" # key alias pattern ``` +**Tag-based** (matches keys/teams by metadata tags, wildcards supported): + +```yaml showLineNumbers title="config.yaml" +policy_attachments: + - policy: hipaa-compliance + tags: + - "healthcare" + - "health-*" # wildcard - matches health-team, health-dev, etc. +``` + +Tags are read from key and team `metadata.tags`. For example, a key created with `metadata: {"tags": ["healthcare"]}` would match the attachment above. + +## Test Policy Matching + +Debug which policies and guardrails apply for a given context. Use this to verify your policy configuration before deploying. + + + + +Go to **Policies** > **Test** tab. Enter a team alias, key alias, model, or tags and click **Test** to see which policies match and what guardrails would be applied. + + + + + + +```bash +curl -X POST "http://localhost:4000/policies/resolve" \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ + "tags": ["healthcare"], + "model": "gpt-4" + }' +``` + +Response: + +```json +{ + "effective_guardrails": ["pii_masking"], + "matched_policies": [ + { + "policy_name": "hipaa-compliance", + "matched_via": "tag:healthcare", + "guardrails_added": ["pii_masking"] + } + ] +} +``` + + + + ## Config Reference ### `policies` @@ -233,14 +341,18 @@ policy_attachments: scope: ... teams: [...] keys: [...] + models: [...] + tags: [...] ``` | Field | Type | Description | |-------|------|-------------| | `policy` | `string` | **Required.** Name of the policy to attach. | | `scope` | `string` | Use `"*"` to apply globally. | -| `teams` | `list[string]` | Team aliases (from `/team/new`). | +| `teams` | `list[string]` | Team aliases (from `/team/new`). Supports `*` wildcard. | | `keys` | `list[string]` | Key aliases (from `/key/generate`). Supports `*` wildcard. | +| `models` | `list[string]` | Model names. Supports `*` wildcard. | +| `tags` | `list[string]` | Tag patterns (from key/team `metadata.tags`). Supports `*` wildcard. | ### Response Headers @@ -248,6 +360,7 @@ policy_attachments: |--------|-------------| | `x-litellm-applied-policies` | Policies that matched this request | | `x-litellm-applied-guardrails` | Guardrails that actually ran | +| `x-litellm-policy-sources` | Why each policy matched (e.g., `hipaa=tag:healthcare; baseline=scope:*`) | ## How it works diff --git a/docs/my-website/docs/proxy/guardrails/policy_tags.md b/docs/my-website/docs/proxy/guardrails/policy_tags.md new file mode 100644 index 00000000000..11840116c31 --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/policy_tags.md @@ -0,0 +1,139 @@ +# Tag-Based Policy Attachments + +Apply guardrail policies automatically to any key or team that has a specific tag. Instead of attaching policies one-by-one, tag your keys and let the policy engine handle the rest. + +**Example:** Your security team requires all healthcare-related keys to run PII masking and PHI detection. Tag those keys with `health`, create a single tag-based attachment, and every matching key gets the guardrails automatically. + +## 1. Create a Policy with Guardrails + +Navigate to **Policies** in the left sidebar. You'll see a list of existing policies along with their guardrails. + +![Policies list page showing existing policies and the + Add New Policy button](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/d7aa1e1f-011e-40bf-a356-6dfe9d5d54f1/ascreenshot_8db95c231a7f4a79a36c2a98ba127542_text_export.jpeg) + +Click **+ Add New Policy**. In the modal, enter a name for your policy (e.g., `high-risk-policy2`). You can also type to search existing policy names if you want to reference them. + +![Create New Policy modal — enter the policy name and optional description](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/18f1ff69-9b83-4a98-9aad-9892a104d3ff/ascreenshot_1c6b85231cad4ec695750b53bbbda52c_text_export.jpeg) + +Scroll down to **Guardrails to Add**. Click the dropdown to see all available guardrails configured on your proxy — select the ones this policy should enforce. + +![Guardrails to Add dropdown showing available guardrails like OAI-moderation, phi-pre-guard, pii-pre-guard](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/55cedad7-9939-44a1-8644-a184cde82ab7/ascreenshot_eab4e55b82b8411893eccb6234d60b82_text_export.jpeg) + +After selecting your guardrails, they appear as chips in the input field. The **Resolved Guardrails** section below shows the final set that will be applied (including any inherited from a parent policy). + +![Selected guardrails shown as chips: testing-pl, phi-pre-guard, pii-pre-guard. Resolved Guardrails preview below.](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/c06d5b08-1c85-4715-b827-3e6864880428/ascreenshot_7a082e55f3ad425f9009346c68afae23_text_export.jpeg) + +Click **Create Policy** to save. + +![Click Create Policy to save the new policy](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/7e6eae64-4bba-4d72-b226-d1308ac576a8/ascreenshot_22d0ed686c594221bbbd2f40df214d75_text_export.jpeg) + +## 2. Add a Tag Attachment for the Policy + +After creating the policy, switch to the **Attachments** tab. This is where you define *where* the policy applies. + +![Switch to the Attachments tab — shows the attachment table and scope documentation](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/871ae6d9-16d1-44e2-baf2-7bb8a9e72087/ascreenshot_76e124619d70462ea0e2fbb46ded1ac9_text_export.jpeg) + +Click **+ Add New Attachment**. The Attachments page explains the available scopes: Global, Teams, Keys, Models, and **Tags**. + +![Attachments page showing scope types including Tags — click + Add New Attachment](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/d45ab8bc-fc1e-425b-8a3f-44d18df810ec/ascreenshot_425824030f3144b7ab3c0ac570349b00_text_export.jpeg) + +In the **Create Policy Attachment** modal, first select the policy you just created from the dropdown. + +![Select the policy to attach from the dropdown (e.g., high-risk-policy2)](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/e0dcac40-e39c-4a6a-9d9c-4bbb9ec0ee91/ascreenshot_445b19894e0b466196a13e20c8e67f2d_text_export.jpeg) + +Choose **Specific (teams, keys, models, or tags)** as the scope type. This expands the form to show fields for Teams, Keys, Models, and Tags. + +![Select "Specific" scope type to reveal the Tags field](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/f685e02a-e22e-4c6c-9742-d5268746214b/ascreenshot_14d63d9d06dd4fc7854cfeb5e8d9ef85_text_export.jpeg) + +Scroll down to the **Tags** field and type the tag to match — here we enter `health`. You can enter any string, or use a wildcard pattern like `health-*` to match all tags starting with `health-` (e.g., `health-team`, `health-dev`). + +![Tags field with "health" entered. Supports wildcards like prod-* matching prod-us, prod-eu.](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/14581df7-732c-4ea5-b36d-58270b00e92c/ascreenshot_e734c81418f046549b61a84b9d352a29_text_export.jpeg) + +## 3. Check the Impact of the Attachment + +Before creating the attachment, click **Estimate Impact** to preview how many keys and teams would be affected. This is your blast-radius check — make sure the scope is what you expect before applying. + +![Click Estimate Impact — the tag "health" is entered and ready to preview](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/6ccb81d7-3d11-48b0-b634-fc4d738aa530/ascreenshot_2eb89e6ff13a4b12b61004660a36c30c_text_export.jpeg) + +The **Impact Preview** appears inline, showing exactly how many keys and teams would be affected. In this example: "This attachment would affect **1 key** and **0 teams**", with the key alias `hi` listed. + +![Impact Preview showing "This attachment would affect 1 key and 0 teams." Keys: hi](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/8834d85a-2c15-48dd-8d6b-810cf11ee5c4/ascreenshot_d814b42ca9f34c23b0c2269bfa3e64fb_text_export.jpeg) + +Once you're satisfied with the impact, click **Create Attachment** to save. + +![Click Create Attachment to finalize](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/4a8918f2-eedb-4f49-a53b-4e46d0387d2a/ascreenshot_b08d490d836d4f46b4e5cbb14f61377a_text_export.jpeg) + +The attachment now appears in the table with the policy name `high-risk-policy2` and tag `health` visible. + +![Attachments table showing the new attachment with policy high-risk-policy2 and tag "health"](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/45867887-0aec-44a4-963b-b6cc6c302e3e/ascreenshot_981caeff98574ec89a8a53cd295e5043_text_export.jpeg) + +## 4. Create a Key with the Tag + +Navigate to **Virtual Keys** in the left sidebar. Click **+ Create New Key**. + +![Virtual Keys page showing existing keys — click + Create New Key](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/4c1f9448-e590-4546-9357-6f68aa395b27/ascreenshot_4a7bc5be9e4347f3a9fe46f78d938d7c_text_export.jpeg) + +Enter a key name and select a model. Then expand **Optional Settings** and scroll down to the **Tags** field. + +![Create New Key modal — enter the key name](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/f84f7a2b-8057-4926-9f80-d68e437c77cf/ascreenshot_a277c8611b6e41059663b0759cd85cab_text_export.jpeg) + +In the **Tags** field, type `health` and press Enter. This is the tag the policy engine will match against. + +![Tags field in key creation — type "health" to add the tag](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/3ad3bf10-76d2-4f15-9a66-ed6c99bb25c4/ascreenshot_8a8773fb65fc49329cb1716da92b2723_text_export.jpeg) + +The tag `health` now appears as a chip in the Tags field. Confirm your settings look correct. + +![Tags field showing "health" selected with a checkmark](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/de3e58a9-6013-4d0c-882e-5517ea286684/ascreenshot_c7eef1736fce4aa894ac3b118b3800a2_text_export.jpeg) + +Click **Create Key** at the bottom of the form. + +![Click Create Key to generate the new virtual key with the health tag](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/51d419ea-ee80-4e24-8e93-b99a844881bc/ascreenshot_097d4564289943a88e30b5d2e3eab262_text_export.jpeg) + +A dialog appears with your new virtual key. Click **Copy Virtual Key** — you'll need this to test in the next step. + +![Save your Key dialog — click Copy Virtual Key to copy it to clipboard](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/e87a0cc1-4d12-4066-bfa2-973159808fd1/ascreenshot_7b616a7291d0497a9c61bdcdb59394d7_text_export.jpeg) + +## 5. Test the Key and Validate the Policy is Applied + +Navigate to **Playground** in the left sidebar to test the key interactively. + +![Navigate to Playground from the sidebar](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/e6f8a3ee-e9e8-4107-93d1-bfca734c5ce9/ascreenshot_539bde38abe646e49148a912fff2d257_text_export.jpeg) + +Under **Virtual Key Source**, select "Virtual Key" and paste the key you just copied into the input field. + +![Paste the virtual key into the Playground configuration](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/a6612c4a-d499-4e54-8019-f54fde674ad9/ascreenshot_e85ebb9051554594bab0da57823fafad_text_export.jpeg) + +Select a model from the **Select Model** dropdown. + +![Select a model (e.g., bedrock-claude-opus-4.5) from the dropdown](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/325e330f-3eff-4c5e-b177-21916138a2f5/ascreenshot_693478f89c034e949e08f3ed0dd05120_text_export.jpeg) + +Type a message and press Enter. If a guardrail blocks the request, you'll see it in the response. In this example, the `testing-pl` guardrail detected an email pattern and returned a 403 error — confirming the policy is working. + +![Guardrail in action — the request was blocked with "Content blocked: email pattern detected"](https://colony-recorder.s3.amazonaws.com/files/2026-02-11/2cf16809-d2e5-4eae-a7dd-6a16dfcca7ce/ascreenshot_727d7d4ed20b4a52b2b41e39fd36eccb_text_export.jpeg) + +**Using curl:** + +You can also verify via the command line. The response headers confirm which policies and guardrails were applied: + +```bash +curl -v http://localhost:4000/chat/completions \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "say hi"}] + }' +``` + +Check the response headers: + +``` +x-litellm-applied-policies: high-risk-policy2 +x-litellm-applied-guardrails: pii-pre-guard,phi-pre-guard,testing-pl +x-litellm-policy-sources: high-risk-policy2=tag:health +``` + +| Header | What it tells you | +|--------|-------------------| +| `x-litellm-applied-policies` | Which policies matched this request | +| `x-litellm-applied-guardrails` | Which guardrails actually ran | +| `x-litellm-policy-sources` | **Why** each policy matched — `tag:health` confirms it was the tag | diff --git a/docs/my-website/docs/tutorials/claude_code_beta_headers.md b/docs/my-website/docs/tutorials/claude_code_beta_headers.md index 9c1645e0277..4cc6f7ff92b 100644 --- a/docs/my-website/docs/tutorials/claude_code_beta_headers.md +++ b/docs/my-website/docs/tutorials/claude_code_beta_headers.md @@ -1,8 +1,8 @@ import Image from '@theme/IdealImage'; -# Claude Code - Fixing Invalid Beta Header Errors +# Claude Code - Managing Anthropic Beta Headers -When using Claude Code with LiteLLM and non-Anthropic providers (Bedrock, Azure AI, Vertex AI), you may encounter "invalid beta header" errors. This guide explains how to fix these errors locally or contribute a fix to LiteLLM. +When using Claude Code with LiteLLM and non-Anthropic providers (Bedrock, Azure AI, Vertex AI), you need to ensure that only supported beta headers are sent to each provider. This guide explains how to add support for new beta headers or fix invalid beta header errors. ## What Are Beta Headers? @@ -12,7 +12,7 @@ Anthropic uses beta headers to enable experimental features in Claude. When you anthropic-beta: prompt-caching-scope-2026-01-05,advanced-tool-use-2025-11-20 ``` -However, not all providers support all Anthropic beta features. When an unsupported beta header is sent to a provider, you'll see an error. +However, not all providers support all Anthropic beta features. LiteLLM uses `anthropic_beta_headers_config.json` to manage which beta headers are supported by each provider. ## Common Error Message @@ -22,17 +22,22 @@ Error: The model returned the following errors: invalid beta flag ## How LiteLLM Handles Beta Headers -LiteLLM automatically filters out unsupported beta headers using a configuration file: +LiteLLM uses a strict validation approach with a configuration file: ``` litellm/litellm/anthropic_beta_headers_config.json ``` -This JSON file lists which beta headers are **unsupported** for each provider. Headers not in the unsupported list are passed through to the provider. +This JSON file contains a **mapping** of beta headers for each provider: +- **Keys**: Input beta header names (from Anthropic) +- **Values**: Provider-specific header names (or `null` if unsupported) +- **Validation**: Only headers present in the mapping with non-null values are forwarded -## Quick Fix: Update Config Locally +This enforces stricter validation than just filtering unsupported headers - headers must be explicitly defined to be allowed. -If you encounter an invalid beta header error, you can fix it immediately by updating the config file locally. +## Adding Support for a New Beta Header + +When Anthropic releases a new beta feature, you need to add it to the configuration file for each provider. ### Step 1: Locate the Config File @@ -46,43 +51,47 @@ cd $(python -c "import litellm; import os; print(os.path.dirname(litellm.__file_ # litellm/anthropic_beta_headers_config.json ``` -### Step 2: Add the Unsupported Header +### Step 2: Add the New Beta Header -Open `anthropic_beta_headers_config.json` and add the problematic header to the appropriate provider's list: +Open `anthropic_beta_headers_config.json` and add the new header to each provider's mapping: ```json title="anthropic_beta_headers_config.json" { - "description": "Unsupported Anthropic beta headers for each provider. Headers listed here will be dropped. Headers not listed are passed through as-is.", - "anthropic": [], - "azure_ai": [], - "bedrock_converse": [ - "prompt-caching-scope-2026-01-05", - "bash_20250124", - "bash_20241022", - "text_editor_20250124", - "text_editor_20241022", - "compact-2026-01-12", - "advanced-tool-use-2025-11-20", - "web-fetch-2025-09-10", - "code-execution-2025-08-25", - "skills-2025-10-02", - "files-api-2025-04-14" - ], - "bedrock": [ - "advanced-tool-use-2025-11-20", - "prompt-caching-scope-2026-01-05", - "structured-outputs-2025-11-13", - "web-fetch-2025-09-10", - "code-execution-2025-08-25", - "skills-2025-10-02", - "files-api-2025-04-14" - ], - "vertex_ai": [ - "prompt-caching-scope-2026-01-05" - ] + "description": "Mapping of Anthropic beta headers for each provider. Keys are input header names, values are provider-specific header names (or null if unsupported). Only headers present in mapping keys with non-null values can be forwarded.", + "anthropic": { + "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", + "new-feature-2026-03-01": "new-feature-2026-03-01", + ... + }, + "azure_ai": { + "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", + "new-feature-2026-03-01": "new-feature-2026-03-01", + ... + }, + "bedrock_converse": { + "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", + "new-feature-2026-03-01": null, + ... + }, + "bedrock": { + "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", + "new-feature-2026-03-01": null, + ... + }, + "vertex_ai": { + "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", + "new-feature-2026-03-01": null, + ... + } } ``` +**Key Points:** +- **Supported headers**: Set the value to the provider-specific header name (often the same as the key) +- **Unsupported headers**: Set the value to `null` +- **Header transformations**: Some providers use different header names (e.g., Bedrock maps `advanced-tool-use-2025-11-20` to `tool-search-tool-2025-10-19`) +- **Alphabetical order**: Keep headers sorted alphabetically for maintainability + ### Step 3: Restart Your Application After updating the config file, restart your LiteLLM proxy or application: @@ -97,9 +106,64 @@ litellm --config config.yaml The updated configuration will be loaded automatically. +## Fixing Invalid Beta Header Errors + +If you encounter an "invalid beta flag" error, it means a beta header is being sent that the provider doesn't support. + +### Step 1: Identify the Problematic Header + +Check your logs to see which header is causing the issue: + +```bash +Error: The model returned the following errors: invalid beta flag: new-feature-2026-03-01 +``` + +### Step 2: Update the Config + +Set the header value to `null` for that provider: + +```json title="anthropic_beta_headers_config.json" +{ + "bedrock_converse": { + "new-feature-2026-03-01": null + } +} +``` + +### Step 3: Restart and Test + +Restart your application and verify the header is now filtered out. + ## Contributing a Fix to LiteLLM -Help the community by contributing your fix! If your local changes work, please raise a PR with the addition of the header and we will merge it. +Help the community by contributing your fix! + +### What to Include in Your PR + +1. **Update the config file**: Add the new beta header to `litellm/anthropic_beta_headers_config.json` +2. **Test your changes**: Verify the header is correctly filtered/mapped for each provider +3. **Documentation**: Include provider documentation links showing which headers are supported + +### Example PR Description + +```markdown +## Add support for new-feature-2026-03-01 beta header + +### Changes +- Added `new-feature-2026-03-01` to anthropic_beta_headers_config.json +- Set to `null` for bedrock_converse (unsupported) +- Set to header name for anthropic, azure_ai (supported) + +### Testing +Tested with: +- ✅ Anthropic: Header passed through correctly +- ✅ Azure AI: Header passed through correctly +- ✅ Bedrock Converse: Header filtered out (returns error without fix) + +### References +- Anthropic docs: [link] +- AWS Bedrock docs: [link] +``` ## How Beta Header Filtering Works @@ -116,14 +180,51 @@ sequenceDiagram CC->>LP: Request with beta headers Note over CC,LP: anthropic-beta: header1,header2,header3 - LP->>Config: Load unsupported headers for provider - Config-->>LP: Returns unsupported list + LP->>Config: Load header mapping for provider + Config-->>LP: Returns mapping (header→value or null) - Note over LP: Filter headers:
- Remove unsupported
- Keep supported + Note over LP: Validate & Transform:
1. Check if header exists in mapping
2. Filter out null values
3. Map to provider-specific names - LP->>Provider: Request with filtered headers - Note over LP,Provider: anthropic-beta: header2
(header1, header3 removed) + LP->>Provider: Request with filtered & mapped headers + Note over LP,Provider: anthropic-beta: mapped-header2
(header1, header3 filtered out) Provider-->>LP: Success response LP-->>CC: Response -``` \ No newline at end of file +``` + +### Filtering Rules + +1. **Header must exist in mapping**: Unknown headers are filtered out +2. **Header must have non-null value**: Headers with `null` values are filtered out +3. **Header transformation**: Headers are mapped to provider-specific names (e.g., `advanced-tool-use-2025-11-20` → `tool-search-tool-2025-10-19` for Bedrock) + +### Example + +Request with headers: +``` +anthropic-beta: advanced-tool-use-2025-11-20,computer-use-2025-01-24,unknown-header +``` + +For Bedrock Converse: +- ✅ `computer-use-2025-01-24` → `computer-use-2025-01-24` (supported, passed through) +- ❌ `advanced-tool-use-2025-11-20` → filtered out (null value in config) +- ❌ `unknown-header` → filtered out (not in config) + +Result sent to Bedrock: +``` +anthropic-beta: computer-use-2025-01-24 +``` + +## Provider-Specific Notes + +### Bedrock +- Beta headers appear in both HTTP headers AND request body (`additionalModelRequestFields.anthropic_beta`) +- Some headers are transformed (e.g., `advanced-tool-use` → `tool-search-tool`) + +### Azure AI +- Uses same header names as Anthropic +- Some features not yet supported (check config for null values) + +### Vertex AI +- Some headers are transformed to match Vertex AI's implementation +- Limited beta feature support compared to Anthropic \ No newline at end of file diff --git a/docs/my-website/img/policy_team_attach.png b/docs/my-website/img/policy_team_attach.png new file mode 100644 index 00000000000..4e337931ed8 Binary files /dev/null and b/docs/my-website/img/policy_team_attach.png differ diff --git a/docs/my-website/img/policy_test_matching.png b/docs/my-website/img/policy_test_matching.png new file mode 100644 index 00000000000..5d024ae78b4 Binary files /dev/null and b/docs/my-website/img/policy_test_matching.png differ diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 2c3dfb2b863..579b5699101 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -42,49 +42,62 @@ const sidebars = { label: "Guardrails", items: [ "proxy/guardrails/quick_start", - "proxy/guardrails/guardrail_policies", "proxy/guardrails/guardrail_load_balancing", + "proxy/guardrails/test_playground", + "proxy/guardrails/litellm_content_filter", { type: "category", - "label": "Contributing to Guardrails", + label: "Providers", + items: [ + ...[ + "proxy/guardrails/qualifire", + "proxy/guardrails/aim_security", + "proxy/guardrails/onyx_security", + "proxy/guardrails/aporia_api", + "proxy/guardrails/azure_content_guardrail", + "proxy/guardrails/bedrock", + "proxy/guardrails/enkryptai", + "proxy/guardrails/ibm_guardrails", + "proxy/guardrails/grayswan", + "proxy/guardrails/hiddenlayer", + "proxy/guardrails/lasso_security", + "proxy/guardrails/guardrails_ai", + "proxy/guardrails/lakera_ai", + "proxy/guardrails/model_armor", + "proxy/guardrails/noma_security", + "proxy/guardrails/dynamoai", + "proxy/guardrails/openai_moderation", + "proxy/guardrails/pangea", + "proxy/guardrails/pillar_security", + "proxy/guardrails/pii_masking_v2", + "proxy/guardrails/panw_prisma_airs", + "proxy/guardrails/secret_detection", + "proxy/guardrails/custom_guardrail", + "proxy/guardrails/custom_code_guardrail", + "proxy/guardrails/prompt_injection", + "proxy/guardrails/tool_permission", + "proxy/guardrails/zscaler_ai_guard", + "proxy/guardrails/javelin" + ].sort(), + ], + }, + { + type: "category", + label: "Contributing to Guardrails", items: [ "adding_provider/generic_guardrail_api", "adding_provider/simple_guardrail_tutorial", "adding_provider/adding_guardrail_support", ] }, - "proxy/guardrails/test_playground", - "proxy/guardrails/litellm_content_filter", - ...[ - "proxy/guardrails/qualifire", - "proxy/guardrails/aim_security", - "proxy/guardrails/onyx_security", - "proxy/guardrails/aporia_api", - "proxy/guardrails/azure_content_guardrail", - "proxy/guardrails/bedrock", - "proxy/guardrails/enkryptai", - "proxy/guardrails/ibm_guardrails", - "proxy/guardrails/grayswan", - "proxy/guardrails/hiddenlayer", - "proxy/guardrails/lasso_security", - "proxy/guardrails/guardrails_ai", - "proxy/guardrails/lakera_ai", - "proxy/guardrails/model_armor", - "proxy/guardrails/noma_security", - "proxy/guardrails/dynamoai", - "proxy/guardrails/openai_moderation", - "proxy/guardrails/pangea", - "proxy/guardrails/pillar_security", - "proxy/guardrails/pii_masking_v2", - "proxy/guardrails/panw_prisma_airs", - "proxy/guardrails/secret_detection", - "proxy/guardrails/custom_guardrail", - "proxy/guardrails/custom_code_guardrail", - "proxy/guardrails/prompt_injection", - "proxy/guardrails/tool_permission", - "proxy/guardrails/zscaler_ai_guard", - "proxy/guardrails/javelin" - ].sort(), + ], + }, + { + type: "category", + label: "Policies", + items: [ + "proxy/guardrails/guardrail_policies", + "proxy/guardrails/policy_tags", ], }, { @@ -396,6 +409,16 @@ const sidebars = { ], }, "proxy/caching", + { + type: "link", + label: "Guardrails", + href: "https://docs.litellm.ai/docs/proxy/guardrails/quick_start", + }, + { + type: "link", + label: "Policies", + href: "https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies", + }, { type: "category", label: "Create Custom Plugins", @@ -1085,6 +1108,17 @@ const sidebars = { "troubleshoot/max_callbacks", ], }, + { + type: "category", + label: "Blog", + items: [ + { + type: "link", + label: "Incident: Broken Model Cost Map", + href: "/blog/model-cost-map-incident", + }, + ], + }, ], }; diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 569ea17f6d8..a41b3f3bf6f 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -899,49 +899,49 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): batch_id=response.id, model_id=model_id ) - if ( - response.output_file_id and model_id - ): # return a file id with the model_id and output_file_id - original_output_file_id = response.output_file_id - response.output_file_id = self.get_unified_output_file_id( - output_file_id=response.output_file_id, - model_id=model_id, - model_name=model_name, - ) - - # Fetch the actual file object for the output file - file_object = None - try: - # Use litellm to retrieve the file object from the provider - from litellm import afile_retrieve - file_object = await afile_retrieve( - custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", - file_id=original_output_file_id + # Handle both output_file_id and error_file_id + for file_attr in ["output_file_id", "error_file_id"]: + file_id_value = getattr(response, file_attr, None) + if file_id_value and model_id: + original_file_id = file_id_value + unified_file_id = self.get_unified_output_file_id( + output_file_id=original_file_id, + model_id=model_id, + model_name=model_name, ) - verbose_logger.debug( - f"Successfully retrieved file object for output_file_id={original_output_file_id}" + setattr(response, file_attr, unified_file_id) + + # Fetch the actual file object from the provider + file_object = None + try: + # Use litellm to retrieve the file object from the provider + from litellm import afile_retrieve + file_object = await afile_retrieve( + custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", + file_id=original_file_id + ) + verbose_logger.debug( + f"Successfully retrieved file object for {file_attr}={original_file_id}" + ) + except Exception as e: + verbose_logger.warning( + f"Failed to retrieve file object for {file_attr}={original_file_id}: {str(e)}. Storing with None and will fetch on-demand." + ) + + await self.store_unified_file_id( + file_id=unified_file_id, + file_object=file_object, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + model_mappings={model_id: original_file_id}, + user_api_key_dict=user_api_key_dict, ) - except Exception as e: - verbose_logger.warning( - f"Failed to retrieve file object for output_file_id={original_output_file_id}: {str(e)}. Storing with None and will fetch on-demand." - ) - - await self.store_unified_file_id( - file_id=response.output_file_id, - file_object=file_object, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - model_mappings={model_id: original_output_file_id}, - user_api_key_dict=user_api_key_dict, - ) - asyncio.create_task( - self.store_unified_object_id( - unified_object_id=response.id, - file_object=response, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - model_object_id=original_response_id, - file_purpose="batch", - user_api_key_dict=user_api_key_dict, - ) + await self.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + model_object_id=original_response_id, + file_purpose="batch", + user_api_key_dict=user_api_key_dict, ) elif isinstance(response, LiteLLMFineTuningJob): ## Check if unified_file_id is in the response @@ -958,15 +958,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): response.id = self.get_unified_generic_response_id( model_id=model_id, generic_response_id=response.id ) - asyncio.create_task( - self.store_unified_object_id( - unified_object_id=response.id, - file_object=response, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - model_object_id=original_response_id, - file_purpose="fine-tune", - user_api_key_dict=user_api_key_dict, - ) + await self.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + model_object_id=original_response_id, + file_purpose="fine-tune", + user_api_key_dict=user_api_key_dict, ) elif isinstance(response, AsyncCursorPage): """ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34-py3-none-any.whl new file mode 100644 index 00000000000..175d84543ec Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34.tar.gz new file mode 100644 index 00000000000..e1fcc0c603f Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b1ca1f71c9e..558dfcc9517 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -914,6 +914,7 @@ model LiteLLM_PolicyAttachmentTable { teams String[] @default([]) // Team aliases or patterns keys String[] @default([]) // Key aliases or patterns models String[] @default([]) // Model names or patterns + tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 8937858bbd9..e0a769a5edf 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.33" +version = "0.4.34" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.33" +version = "0.4.34" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/__init__.py b/litellm/__init__.py index 8174b9d2655..538ee727612 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1155,6 +1155,7 @@ from .exceptions import ( BadRequestError, ImageFetchError, NotFoundError, + PermissionDeniedError, RateLimitError, ServiceUnavailableError, BadGatewayError, @@ -1393,6 +1394,7 @@ if TYPE_CHECKING: from .llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig from .llms.volcengine.responses.transformation import VolcEngineResponsesAPIConfig as VolcEngineResponsesAPIConfig from .llms.manus.responses.transformation import ManusResponsesAPIConfig as ManusResponsesAPIConfig + from .llms.perplexity.responses.transformation import PerplexityResponsesConfig as PerplexityResponsesConfig from .llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig from .llms.openai.chat.o_series_transformation import OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config from .llms.anthropic.skills.transformation import AnthropicSkillsConfig as AnthropicSkillsConfig diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index a01fe9c11db..ebe9af9d85c 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -226,6 +226,7 @@ LLM_CONFIG_NAMES = ( "XAIResponsesAPIConfig", "LiteLLMProxyResponsesAPIConfig", "VolcEngineResponsesAPIConfig", + "PerplexityResponsesConfig", "GoogleAIStudioInteractionsConfig", "OpenAIOSeriesConfig", "AnthropicSkillsConfig", @@ -274,6 +275,7 @@ LLM_CONFIG_NAMES = ( "LmStudioEmbeddingConfig", "NscaleConfig", "PerplexityChatConfig", + "PerplexityResponsesConfig", "AzureOpenAIO1Config", "IBMWatsonXAIConfig", "IBMWatsonXChatConfig", @@ -901,6 +903,10 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.manus.responses.transformation", "ManusResponsesAPIConfig", ), + "PerplexityResponsesConfig": ( + ".llms.perplexity.responses.transformation", + "PerplexityResponsesConfig", + ), "GoogleAIStudioInteractionsConfig": ( ".llms.gemini.interactions.transformation", "GoogleAIStudioInteractionsConfig", diff --git a/litellm/_logging.py b/litellm/_logging.py index e222627e76c..fd833f7056a 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,9 +1,13 @@ -import json +import ast import logging import os import sys from datetime import datetime from logging import Formatter +from typing import Any, Dict, Optional + +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads set_verbose = False @@ -19,6 +23,67 @@ handler = logging.StreamHandler() handler.setLevel(numeric_level) +def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]: + """ + Try to parse a log message as JSON. Returns parsed dict if valid, else None. + Handles messages that are entirely valid JSON (e.g. json.dumps output). + Uses shared safe_json_loads for consistent error handling. + """ + if not message or not isinstance(message, str): + return None + msg_stripped = message.strip() + if not (msg_stripped.startswith("{") or msg_stripped.startswith("[")): + return None + parsed = safe_json_loads(message, default=None) + if parsed is None or not isinstance(parsed, dict): + return None + return parsed + + +def _try_parse_embedded_python_dict(message: str) -> Optional[Dict[str, Any]]: + """ + Try to find and parse a Python dict repr (e.g. str(d) or repr(d)) embedded in + the message. Handles patterns like: + "get_available_deployment for model: X, Selected deployment: {'model_name': '...', ...} for model: X" + Uses ast.literal_eval for safe parsing. Returns the parsed dict or None. + """ + if not message or not isinstance(message, str) or "{" not in message: + return None + i = 0 + while i < len(message): + start = message.find("{", i) + if start == -1: + break + depth = 0 + for j in range(start, len(message)): + c = message[j] + if c == "{": + depth += 1 + elif c == "}": + depth -= 1 + if depth == 0: + substr = message[start : j + 1] + try: + result = ast.literal_eval(substr) + if isinstance(result, dict) and len(result) > 0: + return result + except (ValueError, SyntaxError, TypeError): + pass + break + i = start + 1 + return None + + +# Standard LogRecord attribute names - used to identify 'extra' fields. +# Derived at runtime so we automatically include version-specific attrs (e.g. taskName). +def _get_standard_record_attrs() -> frozenset: + """Standard LogRecord attribute names - excludes extra keys from logger.debug(..., extra={...}).""" + return frozenset(logging.LogRecord("", 0, "", 0, "", (), None).__dict__.keys()) + + +_STANDARD_RECORD_ATTRS = _get_standard_record_attrs() + + class JsonFormatter(Formatter): def __init__(self): super(JsonFormatter, self).__init__() @@ -29,16 +94,31 @@ class JsonFormatter(Formatter): return dt.isoformat() def format(self, record): - json_record = { - "message": record.getMessage(), + message_str = record.getMessage() + json_record: Dict[str, Any] = { + "message": message_str, "level": record.levelname, "timestamp": self.formatTime(record), } + # Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties + parsed = _try_parse_json_message(message_str) + if parsed is None: + parsed = _try_parse_embedded_python_dict(message_str) + if parsed is not None: + for key, value in parsed.items(): + if key not in json_record: + json_record[key] = value + + # Include extra attributes passed via logger.debug("msg", extra={...}) + for key, value in record.__dict__.items(): + if key not in _STANDARD_RECORD_ATTRS and key not in json_record: + json_record[key] = value + if record.exc_info: json_record["stacktrace"] = self.formatException(record.exc_info) - return json.dumps(json_record) + return safe_dumps(json_record) # Function to set up exception handlers for JSON logging @@ -169,15 +249,15 @@ def _initialize_loggers_with_handler(handler: logging.Handler): def _get_uvicorn_json_log_config(): """ Generate a uvicorn log_config dictionary that applies JSON formatting to all loggers. - + This ensures that uvicorn's access logs, error logs, and all application logs are formatted as JSON when json_logs is enabled. """ json_formatter_class = "litellm._logging.JsonFormatter" - + # Use the module-level log_level variable for consistency uvicorn_log_level = log_level.upper() - + log_config = { "version": 1, "disable_existing_loggers": False, @@ -222,7 +302,7 @@ def _get_uvicorn_json_log_config(): }, }, } - + return log_config diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 4ebb5ddb609..5edb8067a08 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -1,33 +1,151 @@ { - "description": "Unsupported Anthropic beta headers for each provider. Headers listed here will be dropped. Headers not listed are passed through as-is.", - "anthropic": [], - "azure_ai": [], - "bedrock_converse": [ - "prompt-caching-scope-2026-01-05", - "bash_20250124", - "bash_20241022", - "text_editor_20250124", - "text_editor_20241022", - "compact-2026-01-12", - "advanced-tool-use-2025-11-20", - "web-fetch-2025-09-10", - "code-execution-2025-08-25", - "skills-2025-10-02", - "files-api-2025-04-14", - "fast-mode-2026-02-01" - ], - "bedrock": [ - "advanced-tool-use-2025-11-20", - "prompt-caching-scope-2026-01-05", - "structured-outputs-2025-11-13", - "web-fetch-2025-09-10", - "code-execution-2025-08-25", - "skills-2025-10-02", - "files-api-2025-04-14", - "fast-mode-2026-02-01", - "mcp-servers-2025-12-04" - ], - "vertex_ai": [ - "prompt-caching-scope-2026-01-05" - ] -} + "description": "Mapping of Anthropic beta headers for each provider. Keys are input header names, values are provider-specific header names (or null if unsupported). Only headers present in mapping keys with non-null values can be forwarded.", + "anthropic": { + "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", + "bash_20241022": "bash_20241022", + "bash_20250124": "bash_20250124", + "code-execution-2025-08-25": "code-execution-2025-08-25", + "compact-2026-01-12": "compact-2026-01-12", + "computer-use-2025-01-24": "computer-use-2025-01-24", + "computer-use-2025-11-24": "computer-use-2025-11-24", + "context-1m-2025-08-07": "context-1m-2025-08-07", + "context-management-2025-06-27": "context-management-2025-06-27", + "effort-2025-11-24": "effort-2025-11-24", + "fast-mode-2026-02-01": "fast-mode-2026-02-01", + "files-api-2025-04-14": "files-api-2025-04-14", + "structured-output-2024-03-01": "structured-output-2024-03-01", + "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", + "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", + "mcp-client-2025-11-20": "mcp-client-2025-11-20", + "mcp-client-2025-04-04": "mcp-client-2025-04-04", + "mcp-servers-2025-12-04": "mcp-servers-2025-12-04", + "output-128k-2025-02-19": "output-128k-2025-02-19", + "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05", + "skills-2025-10-02": "skills-2025-10-02", + "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", + "text_editor_20241022": "text_editor_20241022", + "text_editor_20250124": "text_editor_20250124", + "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19", + "web-fetch-2025-09-10": "web-fetch-2025-09-10", + "web-search-2025-03-05": "web-search-2025-03-05" + }, + "azure_ai": { + "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", + "bash_20241022": "bash_20241022", + "bash_20250124": "bash_20250124", + "code-execution-2025-08-25": "code-execution-2025-08-25", + "compact-2026-01-12": null, + "computer-use-2025-01-24": "computer-use-2025-01-24", + "computer-use-2025-11-24": "computer-use-2025-11-24", + "context-1m-2025-08-07": "context-1m-2025-08-07", + "context-management-2025-06-27": "context-management-2025-06-27", + "effort-2025-11-24": "effort-2025-11-24", + "fast-mode-2026-02-01": null, + "files-api-2025-04-14": "files-api-2025-04-14", + "fine-grained-tool-streaming-2025-05-14": null, + "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", + "mcp-client-2025-11-20": "mcp-client-2025-11-20", + "mcp-client-2025-04-04": "mcp-client-2025-04-04", + "mcp-servers-2025-12-04": "mcp-servers-2025-12-04", + "output-128k-2025-02-19": null, + "structured-output-2024-03-01": null, + "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05", + "skills-2025-10-02": "skills-2025-10-02", + "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", + "text_editor_20241022": null, + "text_editor_20250124": null, + "token-efficient-tools-2025-02-19": null, + "web-fetch-2025-09-10": "web-fetch-2025-09-10", + "web-search-2025-03-05": "web-search-2025-03-05" + }, + "bedrock_converse": { + "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", + "bash_20241022": null, + "bash_20250124": null, + "code-execution-2025-08-25": null, + "compact-2026-01-12": null, + "computer-use-2025-01-24": "computer-use-2025-01-24", + "computer-use-2025-11-24": "computer-use-2025-11-24", + "context-1m-2025-08-07": null, + "context-management-2025-06-27": "context-management-2025-06-27", + "effort-2025-11-24": null, + "fast-mode-2026-02-01": null, + "files-api-2025-04-14": null, + "fine-grained-tool-streaming-2025-05-14": null, + "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", + "mcp-client-2025-11-20": null, + "mcp-client-2025-04-04": null, + "mcp-servers-2025-12-04": null, + "output-128k-2025-02-19": null, + "structured-output-2024-03-01": null, + "prompt-caching-scope-2026-01-05": null, + "skills-2025-10-02": null, + "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", + "text_editor_20241022": null, + "text_editor_20250124": null, + "token-efficient-tools-2025-02-19": null, + "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", + "web-fetch-2025-09-10": null, + "web-search-2025-03-05": null + }, + "bedrock": { + "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", + "bash_20241022": null, + "bash_20250124": null, + "code-execution-2025-08-25": null, + "compact-2026-01-12": "compact-2026-01-12", + "computer-use-2025-01-24": "computer-use-2025-01-24", + "computer-use-2025-11-24": "computer-use-2025-11-24", + "context-1m-2025-08-07": "context-1m-2025-08-07", + "context-management-2025-06-27": "context-management-2025-06-27", + "effort-2025-11-24": null, + "fast-mode-2026-02-01": null, + "files-api-2025-04-14": null, + "fine-grained-tool-streaming-2025-05-14": null, + "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", + "mcp-client-2025-11-20": null, + "mcp-client-2025-04-04": null, + "mcp-servers-2025-12-04": null, + "output-128k-2025-02-19": null, + "structured-output-2024-03-01": null, + "prompt-caching-scope-2026-01-05": null, + "skills-2025-10-02": null, + "structured-outputs-2025-11-13": null, + "text_editor_20241022": null, + "text_editor_20250124": null, + "token-efficient-tools-2025-02-19": null, + "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", + "web-fetch-2025-09-10": null, + "web-search-2025-03-05": null + }, + "vertex_ai": { + "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", + "bash_20241022": null, + "bash_20250124": null, + "code-execution-2025-08-25": null, + "compact-2026-01-12": null, + "computer-use-2025-01-24": "computer-use-2025-01-24", + "computer-use-2025-11-24": "computer-use-2025-11-24", + "context-1m-2025-08-07": null, + "context-management-2025-06-27": "context-management-2025-06-27", + "effort-2025-11-24": null, + "fast-mode-2026-02-01": null, + "files-api-2025-04-14": null, + "fine-grained-tool-streaming-2025-05-14": null, + "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", + "mcp-client-2025-11-20": null, + "mcp-client-2025-04-04": null, + "mcp-servers-2025-12-04": null, + "output-128k-2025-02-19": null, + "structured-output-2024-03-01": null, + "prompt-caching-scope-2026-01-05": null, + "skills-2025-10-02": null, + "structured-outputs-2025-11-13": null, + "text_editor_20241022": null, + "text_editor_20250124": null, + "token-efficient-tools-2025-02-19": null, + "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", + "web-fetch-2025-09-10": null, + "web-search-2025-03-05": "web-search-2025-03-05" + } +} \ No newline at end of file diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 2643f4c03fa..9730ae02698 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -2,14 +2,15 @@ Centralized manager for Anthropic beta headers across different providers. This module provides utilities to: -1. Load beta header configuration from JSON (lists unsupported headers per provider) -2. Filter out unsupported beta headers +1. Load beta header configuration from JSON (mapping of supported headers per provider) +2. Filter and map beta headers based on provider support 3. Handle provider-specific header name mappings (e.g., advanced-tool-use -> tool-search-tool) Design: -- JSON config lists UNSUPPORTED headers for each provider -- Headers not in the unsupported list are passed through -- Header mappings allow renaming headers for specific providers +- JSON config contains mapping of beta headers for each provider +- Keys are input header names, values are provider-specific header names (or null if unsupported) +- Only headers present in mapping keys with non-null values can be forwarded +- This enforces stricter validation than the previous unsupported list approach """ import json @@ -47,13 +48,13 @@ def _load_beta_headers_config() -> Dict: return _BETA_HEADERS_CONFIG except Exception as e: verbose_logger.error(f"Failed to load beta headers config: {e}") - # Return empty config as fallback + # Return empty config as fallback (empty mappings) return { - "anthropic": [], - "azure_ai": [], - "bedrock": [], - "bedrock_converse": [], - "vertex_ai": [] + "anthropic": {}, + "azure_ai": {}, + "bedrock": {}, + "bedrock_converse": {}, + "vertex_ai": {} } @@ -77,21 +78,19 @@ def filter_and_transform_beta_headers( provider: str, ) -> List[str]: """ - Filter beta headers based on provider's unsupported list. + Filter and transform beta headers based on provider's mapping configuration. This function: - 1. Removes headers that are in the provider's unsupported list - 2. Passes through all other headers as-is - - Note: Header transformations/mappings (e.g., advanced-tool-use -> tool-search-tool) - are handled in each provider's transformation code, not here. + 1. Only allows headers that are present in the provider's mapping keys + 2. Filters out headers with null values (unsupported) + 3. Maps headers to provider-specific names (e.g., advanced-tool-use -> tool-search-tool) Args: beta_headers: List of Anthropic beta header values provider: Provider name (e.g., "anthropic", "bedrock", "vertex_ai") Returns: - List of filtered beta headers for the provider + List of filtered and transformed beta headers for the provider """ if not beta_headers: return [] @@ -99,23 +98,33 @@ def filter_and_transform_beta_headers( config = _load_beta_headers_config() provider = get_provider_name(provider) - # Get unsupported headers for this provider - unsupported_headers = set(config.get(provider, [])) + # Get the header mapping for this provider + provider_mapping = config.get(provider, {}) filtered_headers: Set[str] = set() for header in beta_headers: header = header.strip() - # Skip if header is unsupported - if header in unsupported_headers: + # Check if header is in the mapping + if header not in provider_mapping: + verbose_logger.debug( + f"Dropping unknown beta header '{header}' for provider '{provider}' (not in mapping)" + ) + continue + + # Get the mapped header value + mapped_header = provider_mapping[header] + + # Skip if header is unsupported (null value) + if mapped_header is None: verbose_logger.debug( f"Dropping unsupported beta header '{header}' for provider '{provider}'" ) continue - # Pass through as-is - filtered_headers.add(header) + # Add the mapped header + filtered_headers.add(mapped_header) return sorted(list(filtered_headers)) @@ -132,12 +141,14 @@ def is_beta_header_supported( provider: Provider name Returns: - True if the header is supported (not in unsupported list), False otherwise + True if the header is in the mapping with a non-null value, False otherwise """ config = _load_beta_headers_config() provider = get_provider_name(provider) - unsupported_headers = set(config.get(provider, [])) - return beta_header not in unsupported_headers + provider_mapping = config.get(provider, {}) + + # Header is supported if it's in the mapping and has a non-null value + return beta_header in provider_mapping and provider_mapping[beta_header] is not None def get_provider_beta_header( @@ -145,27 +156,29 @@ def get_provider_beta_header( provider: str, ) -> Optional[str]: """ - Check if a beta header is supported by a provider. + Get the provider-specific beta header name for a given Anthropic beta header. - Note: This does NOT handle header transformations/mappings. - Those are handled in each provider's transformation code. + This function handles header transformations/mappings (e.g., advanced-tool-use -> tool-search-tool). Args: anthropic_beta_header: The Anthropic beta header value provider: Provider name Returns: - The original header if supported, or None if unsupported + The provider-specific header name if supported, or None if unsupported/unknown """ config = _load_beta_headers_config() provider = get_provider_name(provider) - # Check if unsupported - unsupported_headers = set(config.get(provider, [])) - if anthropic_beta_header in unsupported_headers: + # Get the header mapping for this provider + provider_mapping = config.get(provider, {}) + + # Check if header is in the mapping + if anthropic_beta_header not in provider_mapping: return None - return anthropic_beta_header + # Return the mapped value (could be None if unsupported) + return provider_mapping[anthropic_beta_header] def update_headers_with_filtered_beta( @@ -208,7 +221,7 @@ def update_headers_with_filtered_beta( def get_unsupported_headers(provider: str) -> List[str]: """ - Get all beta headers that are unsupported by a provider. + Get all beta headers that are unsupported by a provider (have null values in mapping). Args: provider: Provider name @@ -218,4 +231,7 @@ def get_unsupported_headers(provider: str) -> List[str]: """ config = _load_beta_headers_config() provider = get_provider_name(provider) - return config.get(provider, []) + provider_mapping = config.get(provider, {}) + + # Return headers with null values + return [header for header, value in provider_mapping.items() if value is None] diff --git a/litellm/batch_completion/main.py b/litellm/batch_completion/main.py index 7100fb004f8..446e3f2f990 100644 --- a/litellm/batch_completion/main.py +++ b/litellm/batch_completion/main.py @@ -237,17 +237,37 @@ def batch_completion_models_all_responses(*args, **kwargs): if "model" in kwargs: kwargs.pop("model") if "models" in kwargs: - models = kwargs["models"] - kwargs.pop("models") + models = kwargs.pop("models") else: raise Exception("'models' param not in kwargs") + if isinstance(models, str): + models = [models] + elif isinstance(models, (list, tuple)): + models = list(models) + else: + raise TypeError("'models' must be a string or list of strings") + + if len(models) == 0: + return [] + responses = [] with concurrent.futures.ThreadPoolExecutor(max_workers=len(models)) as executor: - for idx, model in enumerate(models): - future = executor.submit(litellm.completion, *args, model=model, **kwargs) - if future.result() is not None: - responses.append(future.result()) + futures = [ + executor.submit(litellm.completion, *args, model=model, **kwargs) + for model in models + ] + + for future in futures: + try: + result = future.result() + if result is not None: + responses.append(result) + except Exception as e: + print_verbose( + f"batch_completion_models_all_responses: model request failed: {str(e)}" + ) + continue return responses diff --git a/litellm/constants.py b/litellm/constants.py index 9c25cf77906..88c57d3ce4c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -48,6 +48,14 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int( os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1) ) DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) + +# Model cost map validation constants +MODEL_COST_MAP_MIN_MODEL_COUNT = int( + os.getenv("MODEL_COST_MAP_MIN_MODEL_COUNT", 50) +) # Minimum number of models a fetched cost map must contain to be considered valid +MODEL_COST_MAP_MAX_SHRINK_RATIO = float( + os.getenv("MODEL_COST_MAP_MAX_SHRINK_RATIO", 0.5) +) # Maximum allowed shrinkage ratio vs local backup (0.5 = reject if fetched map is <50% of backup) DEFAULT_IMAGE_WIDTH = int(os.getenv("DEFAULT_IMAGE_WIDTH", 300)) DEFAULT_IMAGE_HEIGHT = int(os.getenv("DEFAULT_IMAGE_HEIGHT", 300)) # Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit) @@ -205,6 +213,10 @@ REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_bu REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", 2000)) +# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth +LITELLM_ASYNCIO_QUEUE_MAXSIZE = int( + os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000) +) MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int( os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000) ) @@ -1298,6 +1310,9 @@ DEFAULT_SLACK_ALERTING_THRESHOLD = int( os.getenv("DEFAULT_SLACK_ALERTING_THRESHOLD", 300) ) MAX_TEAM_LIST_LIMIT = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20)) +MAX_POLICY_ESTIMATE_IMPACT_ROWS = int( + os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000) +) DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float( os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7) ) diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 9c2f0d95d4d..fe2f9f41f1b 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -28,6 +28,41 @@ else: class ArizeLogger(OpenTelemetry): + """ + Arize logger that sends traces to an Arize endpoint. + + Creates its own dedicated TracerProvider so it can coexist with the + generic ``otel`` callback (or any other OTEL-based integration) without + fighting over the global ``opentelemetry.trace`` TracerProvider singleton. + """ + + def _init_tracing(self, tracer_provider): + """ + Override to always create a *private* TracerProvider for Arize. + + See ArizePhoenixLogger._init_tracing for full rationale. + """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import SpanKind + + if tracer_provider is not None: + self.tracer = tracer_provider.get_tracer("litellm") + self.span_kind = SpanKind + return + + provider = TracerProvider(resource=self._get_litellm_resource(self.config)) + provider.add_span_processor(self._get_span_processor()) + self.tracer = provider.get_tracer("litellm") + self.span_kind = SpanKind + + def _init_otel_logger_on_litellm_proxy(self): + """ + Override: Arize should NOT overwrite the proxy's + ``open_telemetry_logger``. That attribute is reserved for the + primary ``otel`` callback which handles proxy-level parent spans. + """ + pass + def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]): ArizeLogger.set_arize_attributes(span, kwargs, response_obj) return diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index cd345a7f76d..1b038c098f8 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -5,43 +5,211 @@ from litellm._logging import verbose_logger from litellm.integrations.arize import _utils from litellm.integrations.arize._utils import ArizeOTELAttributes from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig -from litellm.integrations.opentelemetry import OpenTelemetry if TYPE_CHECKING: + from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import Span as _Span + from opentelemetry.trace import SpanKind + from litellm.integrations.opentelemetry import OpenTelemetry as _OpenTelemetry from litellm.integrations.opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig from litellm.types.integrations.arize import Protocol as _Protocol Protocol = _Protocol OpenTelemetryConfig = _OpenTelemetryConfig Span = Union[_Span, Any] + OpenTelemetry = _OpenTelemetry else: Protocol = Any OpenTelemetryConfig = Any Span = Any + TracerProvider = Any + SpanKind = Any + # Import OpenTelemetry at runtime + try: + from litellm.integrations.opentelemetry import OpenTelemetry + except ImportError: + OpenTelemetry = None # type: ignore ARIZE_HOSTED_PHOENIX_ENDPOINT = "https://otlp.arize.com/v1/traces" -class ArizePhoenixLogger(OpenTelemetry): +class ArizePhoenixLogger(OpenTelemetry): # type: ignore + """ + Arize Phoenix logger that sends traces to a Phoenix endpoint. + + Creates its own dedicated TracerProvider so it can coexist with the + generic ``otel`` callback (or any other OTEL-based integration) without + fighting over the global ``opentelemetry.trace`` TracerProvider singleton. + """ + + def _init_tracing(self, tracer_provider): + """ + Override to always create a *private* TracerProvider for Arize Phoenix. + + The base ``OpenTelemetry._init_tracing`` falls back to the global + TracerProvider when one already exists. That causes whichever + integration initialises second to silently reuse the first one's + exporter, so spans only reach one destination. + + By creating our own provider we guarantee Arize Phoenix always gets + its own exporter pipeline, regardless of initialisation order. + """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import SpanKind + + if tracer_provider is not None: + # Explicitly supplied (e.g. in tests) — honour it. + self.tracer = tracer_provider.get_tracer("litellm") + self.span_kind = SpanKind + return + + # Always create a dedicated provider — never touch the global one. + provider = TracerProvider(resource=self._get_litellm_resource(self.config)) + provider.add_span_processor(self._get_span_processor()) + self.tracer = provider.get_tracer("litellm") + self.span_kind = SpanKind + verbose_logger.debug( + "ArizePhoenixLogger: Created dedicated TracerProvider " + "(endpoint=%s, exporter=%s)", + self.config.endpoint, + self.config.exporter, + ) + + def _init_otel_logger_on_litellm_proxy(self): + """ + Override: Arize Phoenix should NOT overwrite the proxy's + ``open_telemetry_logger``. That attribute is reserved for the + primary ``otel`` callback which handles proxy-level parent spans. + """ + pass + def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]): ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj) return @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): + from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute + _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) - - # Set project name on the span for all traces to go to custom Phoenix projects - config = ArizePhoenixLogger.get_arize_phoenix_config() - if config.project_name: - from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute - safe_set_attribute(span, "openinference.project.name", config.project_name) - + + # Dynamic project name: check metadata first, then fall back to env var config + dynamic_project_name = ArizePhoenixLogger._get_dynamic_project_name(kwargs) + if dynamic_project_name: + safe_set_attribute(span, "openinference.project.name", dynamic_project_name) + else: + # Fall back to static config from env var + config = ArizePhoenixLogger.get_arize_phoenix_config() + if config.project_name: + safe_set_attribute(span, "openinference.project.name", config.project_name) + return + @staticmethod + def _get_dynamic_project_name(kwargs) -> Optional[str]: + """ + Retrieve dynamic Phoenix project name from request metadata. + + Users can set `metadata.phoenix_project_name` in their request to route + traces to different Phoenix projects dynamically. + """ + standard_logging_payload = kwargs.get("standard_logging_object") + if isinstance(standard_logging_payload, dict): + metadata = standard_logging_payload.get("metadata") + if isinstance(metadata, dict): + project_name = metadata.get("phoenix_project_name") + if project_name: + return str(project_name) + + # Also check litellm_params.metadata for SDK usage + litellm_params = kwargs.get("litellm_params") + if isinstance(litellm_params, dict): + metadata = litellm_params.get("metadata") or {} + else: + metadata = {} + if isinstance(metadata, dict): + project_name = metadata.get("phoenix_project_name") + if project_name: + return str(project_name) + + return None + + def _handle_success(self, kwargs, response_obj, start_time, end_time): + """ + Override to prevent creating duplicate litellm_request spans when a proxy parent span exists. + + ArizePhoenixLogger should reuse the proxy parent span instead of creating a new litellm_request span, + to maintain a shallow span hierarchy as expected by Arize Phoenix. + """ + from opentelemetry.trace import Status, StatusCode + from litellm.secret_managers.main import get_secret_bool + from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME + + verbose_logger.debug( + "ArizePhoenixLogger: Logging kwargs: %s, OTEL config settings=%s", + kwargs, + self.config, + ) + ctx, parent_span = self._get_span_context(kwargs) + + # ArizePhoenixLogger NEVER creates a litellm_request span when a proxy parent span exists + # This is different from the base OpenTelemetry behavior which respects USE_OTEL_LITELLM_REQUEST_SPAN + should_create_primary_span = parent_span is None or ( + parent_span.name != LITELLM_PROXY_REQUEST_SPAN_NAME + and get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN") + ) + + if should_create_primary_span: + # Create a new litellm_request span + span = self._start_primary_span( + kwargs, response_obj, start_time, end_time, ctx + ) + # Raw-request sub-span (if enabled) - child of litellm_request span + self._maybe_log_raw_request( + kwargs, response_obj, start_time, end_time, span + ) + # Ensure proxy-request parent span is annotated with the actual operation kind + if ( + parent_span is not None + and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME + ): + self.set_attributes(parent_span, kwargs, response_obj) + else: + # Do not create primary span (keep hierarchy shallow when parent exists) + span = None + # Only set attributes if the span is still recording (not closed) + # Note: parent_span is guaranteed to be not None here + if parent_span.is_recording(): + parent_span.set_status(Status(StatusCode.OK)) + self.set_attributes(parent_span, kwargs, response_obj) + # Raw-request as direct child of parent_span + self._maybe_log_raw_request( + kwargs, response_obj, start_time, end_time, parent_span + ) + + # 3. Guardrail span + self._create_guardrail_span(kwargs=kwargs, context=ctx) + + # 4. Metrics & cost recording + self._record_metrics(kwargs, response_obj, start_time, end_time) + + # 5. Semantic logs. + if self.config.enable_events: + log_span = span if span is not None else parent_span + if log_span is not None: + self._emit_semantic_logs(kwargs, response_obj, log_span) + + # 6. Do NOT end parent span - it should be managed by its creator + # External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM + # However, proxy-created spans should be closed here + if ( + parent_span is not None + and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME + ): + parent_span.end(end_time=self._to_ns(end_time)) + @staticmethod def get_arize_phoenix_config() -> ArizePhoenixConfig: """ diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index e06b944a419..c36833a6dbf 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -103,10 +103,15 @@ class CBFTransformer: # Use team_alias if available, otherwise team_id, otherwise fallback to 'unknown' entity_id = str(team_alias) if team_alias else (str(team_id) if team_id else 'unknown') + # Get alias fields if they exist + api_key_alias = row.get('api_key_alias') + organization_alias = row.get('organization_alias') + project_alias = row.get('project_alias') + user_alias = row.get('user_alias') + dimensions = { 'entity_type': CZEntityType.TEAM.value, 'entity_id': entity_id, - 'team_id': str(team_id) if team_id else 'unknown', 'team_alias': str(team_alias) if team_alias else 'unknown', 'model': model, 'model_group': str(row.get('model_group', '')), @@ -119,28 +124,37 @@ class CBFTransformer: 'failed_requests': str(row.get('failed_requests', 0)), 'cache_creation_tokens': str(row.get('cache_creation_input_tokens', 0)), 'cache_read_tokens': str(row.get('cache_read_input_tokens', 0)), + 'organization_alias': str(organization_alias) if organization_alias else '', + 'project_alias': str(project_alias) if project_alias else '', + 'user_alias': str(user_alias) if user_alias else '', } # Extract CZRN components to populate corresponding CBF columns czrn_components = self.czrn_generator.extract_components(resource_id) service_type, provider, region, owner_account_id, resource_type, cloud_local_id = czrn_components + # Build resource/account as concat of api_key_alias and api_key_prefix + resource_account = f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash + # CloudZero CBF format with proper column names cbf_record = { # Required CBF fields 'time/usage_start': usage_date.isoformat() if usage_date else None, # Required: ISO-formatted UTC datetime 'cost/cost': float(row.get('spend', 0.0)), # Required: billed cost - 'resource/id': resource_id, # Required when resource tags are present + 'resource/id': model, # Send model name # Usage metrics for token consumption 'usage/amount': total_tokens, # Numeric value of tokens consumed 'usage/units': 'tokens', # Description of token units - # CBF fields that correspond to CZRN components - 'resource/service': service_type, # Maps to CZRN service-type (litellm) - 'resource/account': owner_account_id, # Maps to CZRN owner-account-id (entity_id) + # CBF fields - updated per LIT-1907 + 'resource/service': str(row.get('model_group', '')), # Send model_group + 'resource/account': resource_account, # Send api_key_alias|api_key_prefix 'resource/region': region, # Maps to CZRN region (cross-region) - 'resource/usage_family': resource_type, # Maps to CZRN resource-type (llm-usage) + 'resource/usage_family': str(row.get('custom_llm_provider', '')), # Send provider + + # Action field + 'action/operation': str(team_id) if team_id else '', # Send team_id # Line item details 'lineitem/type': 'Usage', # Standard usage line item @@ -155,13 +169,11 @@ class CBFTransformer: if value and value != 'N/A' and value != 'unknown': # Only add meaningful tags cbf_record[f'resource/tag:{key}'] = str(value) - # Add token breakdown as resource tags for analysis + # Add token breakdown as resource tags for analysis (excluding total_tokens per LIT-1907) if prompt_tokens > 0: cbf_record['resource/tag:prompt_tokens'] = str(prompt_tokens) if completion_tokens > 0: cbf_record['resource/tag:completion_tokens'] = str(completion_tokens) - if total_tokens > 0: - cbf_record['resource/tag:total_tokens'] = str(total_tokens) return CBFRecord(cbf_record) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index bbd55a59bce..407bc581f71 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -616,6 +616,7 @@ class CustomGuardrail(CustomLogger): end_time: Optional[float] = None, duration: Optional[float] = None, event_type: Optional[GuardrailEventHooks] = None, + original_inputs: Optional[Dict] = None, ): """ Add StandardLoggingGuardrailInformation to the request data @@ -625,6 +626,17 @@ class CustomGuardrail(CustomLogger): # Convert None to empty dict to satisfy type requirements guardrail_response = {} if response is None else response + # For apply_guardrail functions in custom_code_guardrail scenario, + # simplify the logged response to "allow", "deny", or "mask" + if original_inputs is not None and isinstance(response, dict): + # Check if inputs were modified by comparing them + if self._inputs_were_modified(original_inputs, response): + guardrail_response = "mask" + else: + guardrail_response = "allow" + + verbose_logger.debug(f"Guardrail response: {response}") + self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, request_data=request_data, @@ -650,8 +662,14 @@ class CustomGuardrail(CustomLogger): This gets logged on downsteam Langfuse, DataDog, etc. """ + # For custom_code_guardrail scenario, log as "deny" instead of full exception + # Check if this is from custom_code_guardrail by checking the class name + guardrail_response: Union[Exception, str] = e + if "CustomCodeGuardrail" in self.__class__.__name__: + guardrail_response = "deny" + self.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=e, + guardrail_json_response=guardrail_response, request_data=request_data, guardrail_status="guardrail_failed_to_respond", duration=duration, @@ -661,6 +679,25 @@ class CustomGuardrail(CustomLogger): ) raise e + def _inputs_were_modified(self, original_inputs: Dict, response: Dict) -> bool: + """ + Compare original inputs with response to determine if content was modified. + + Returns True if the inputs were modified (mask scenario), False otherwise (allow scenario). + """ + # Get all keys from both dictionaries + all_keys = set(original_inputs.keys()) | set(response.keys()) + + # Compare each key's value + for key in all_keys: + original_value = original_inputs.get(key) + response_value = response.get(key) + if original_value != response_value: + return True + + # No modifications detected + return False + def mask_content_in_string( self, content_string: str, @@ -768,6 +805,12 @@ def log_guardrail_information(func): self: CustomGuardrail = args[0] request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} event_type = _infer_event_type_from_function_name(func.__name__) + + # Store original inputs for comparison (for apply_guardrail functions) + original_inputs = None + if func.__name__ == "apply_guardrail" and "inputs" in kwargs: + original_inputs = kwargs.get("inputs") + try: response = await func(*args, **kwargs) return self._process_response( @@ -777,6 +820,7 @@ def log_guardrail_information(func): end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), event_type=event_type, + original_inputs=original_inputs, ) except Exception as e: return self._process_error( @@ -794,6 +838,12 @@ def log_guardrail_information(func): self: CustomGuardrail = args[0] request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} event_type = _infer_event_type_from_function_name(func.__name__) + + # Store original inputs for comparison (for apply_guardrail functions) + original_inputs = None + if func.__name__ == "apply_guardrail" and "inputs" in kwargs: + original_inputs = kwargs.get("inputs") + try: response = func(*args, **kwargs) return self._process_response( @@ -801,6 +851,7 @@ def log_guardrail_information(func): request_data=request_data, duration=(datetime.now() - start_time).total_seconds(), event_type=event_type, + original_inputs=original_inputs, ) except Exception as e: return self._process_error( diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 3cb62905531..0f1ba4a4093 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from urllib.parse import quote from litellm._logging import verbose_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.integrations.additional_logging_utils import AdditionalLoggingUtils from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase from litellm.proxy._types import CommonProxyErrors @@ -41,7 +42,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): batch_size=self.batch_size, flush_interval=self.flush_interval, ) - self.log_queue: asyncio.Queue[GCSLogQueueItem] = asyncio.Queue() # type: ignore[assignment] + self.log_queue: asyncio.Queue[GCSLogQueueItem] = asyncio.Queue( # type: ignore[assignment] + maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE + ) asyncio.create_task(self.periodic_flush()) AdditionalLoggingUtils.__init__(self) @@ -69,6 +72,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") + # When queue is at maxsize, flush immediately to make room (no blocking, no data dropped) + if self.log_queue.full(): + await self.flush_queue() await self.log_queue.put( GCSLogQueueItem( payload=logging_payload, kwargs=kwargs, response_obj=response_obj @@ -91,9 +97,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") - # Add to logging queue - this will be flushed periodically - # Use asyncio.Queue.put() for thread-safe concurrent access - # If queue is full, this will block until space is available (backpressure) + # When queue is at maxsize, flush immediately to make room (no blocking, no data dropped) + if self.log_queue.full(): + await self.flush_queue() await self.log_queue.put( GCSLogQueueItem( payload=logging_payload, kwargs=kwargs, response_obj=response_obj diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 296a88f9a0b..b847180174a 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -5,6 +5,10 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast import litellm from litellm._logging import verbose_logger +from litellm.integrations._types.open_inference import ( + OpenInferenceSpanKindValues, + SpanAttributes, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.secret_managers.main import get_secret_bool @@ -17,10 +21,6 @@ from litellm.types.utils import ( StandardCallbackDynamicParams, StandardLoggingPayload, ) -from litellm.integrations._types.open_inference import ( - OpenInferenceSpanKindValues, - SpanAttributes, -) # OpenTelemetry imports moved to individual functions to avoid import errors when not installed @@ -40,7 +40,9 @@ if TYPE_CHECKING: Context = Union[_Context, Any] SpanExporter = Union[_SpanExporter, Any] UserAPIKeyAuth = Union[_UserAPIKeyAuth, Any] - ManagementEndpointLoggingPayload = Union[_ManagementEndpointLoggingPayload, Any] + ManagementEndpointLoggingPayload = Union[ + _ManagementEndpointLoggingPayload, Any + ] else: Span = Any Tracer = Any @@ -70,6 +72,13 @@ class OpenTelemetryConfig: model_id: Optional[str] = None def __post_init__(self) -> None: + # If endpoint is specified but exporter is still the default "console", + # automatically infer "otlp_http" to send traces to the endpoint. + # This fixes an issue where UI-configured OTEL settings would default + # to console output instead of sending traces to the configured endpoint. + if self.endpoint and isinstance(self.exporter, str) and self.exporter == "console": + self.exporter = "otlp_http" + if not self.service_name: self.service_name = os.getenv("OTEL_SERVICE_NAME", "litellm") if not self.deployment_environment: @@ -95,12 +104,16 @@ class OpenTelemetryConfig: exporter = os.getenv( "OTEL_EXPORTER_OTLP_PROTOCOL", os.getenv("OTEL_EXPORTER", "console") ) - endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", os.getenv("OTEL_ENDPOINT")) + endpoint = os.getenv( + "OTEL_EXPORTER_OTLP_ENDPOINT", os.getenv("OTEL_ENDPOINT") + ) headers = os.getenv( "OTEL_EXPORTER_OTLP_HEADERS", os.getenv("OTEL_HEADERS") ) # example: OTEL_HEADERS=x-honeycomb-team=B85YgLm96***" enable_metrics: bool = ( - os.getenv("LITELLM_OTEL_INTEGRATION_ENABLE_METRICS", "false").lower() + os.getenv( + "LITELLM_OTEL_INTEGRATION_ENABLE_METRICS", "false" + ).lower() == "true" ) enable_events: bool = ( @@ -108,7 +121,9 @@ class OpenTelemetryConfig: == "true" ) service_name = os.getenv("OTEL_SERVICE_NAME", "litellm") - deployment_environment = os.getenv("OTEL_ENVIRONMENT_NAME", "production") + deployment_environment = os.getenv( + "OTEL_ENVIRONMENT_NAME", "production" + ) model_id = os.getenv("OTEL_MODEL_ID", service_name) if exporter == "in_memory": @@ -157,7 +172,9 @@ class OpenTelemetry(CustomLogger): logging.getLogger(__name__) # Enable OpenTelemetry logging - otel_exporter_logger = logging.getLogger("opentelemetry.sdk.trace.export") + otel_exporter_logger = logging.getLogger( + "opentelemetry.sdk.trace.export" + ) otel_exporter_logger.setLevel(logging.DEBUG) # init CustomLogger params @@ -253,7 +270,9 @@ class OpenTelemetry(CustomLogger): # Don't call set_provider to preserve existing context else: # Default proxy provider or unknown type, create our own - verbose_logger.debug("OpenTelemetry: Creating new %s", provider_name) + verbose_logger.debug( + "OpenTelemetry: Creating new %s", provider_name + ) provider = create_new_provider_fn() set_provider_fn(provider) except Exception as e: @@ -274,7 +293,9 @@ class OpenTelemetry(CustomLogger): from opentelemetry.trace import SpanKind def create_tracer_provider(): - provider = TracerProvider(resource=self._get_litellm_resource(self.config)) + provider = TracerProvider( + resource=self._get_litellm_resource(self.config) + ) provider.add_span_processor(self._get_span_processor()) return provider @@ -388,10 +409,14 @@ class OpenTelemetry(CustomLogger): def log_failure_event(self, kwargs, response_obj, start_time, end_time): self._handle_failure(kwargs, response_obj, start_time, end_time) - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event( + self, kwargs, response_obj, start_time, end_time + ): self._handle_success(kwargs, response_obj, start_time, end_time) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event( + self, kwargs, response_obj, start_time, end_time + ): self._handle_failure(kwargs, response_obj, start_time, end_time) async def async_service_success_hook( @@ -588,7 +613,9 @@ class OpenTelemetry(CustomLogger): if dynamic_headers is not None: # Create spans using a temporary tracer with dynamic headers - tracer_to_use = self._get_tracer_with_dynamic_headers(dynamic_headers) + tracer_to_use = self._get_tracer_with_dynamic_headers( + dynamic_headers + ) verbose_logger.debug( "Using dynamic headers for this request: %s", dynamic_headers ) @@ -624,7 +651,9 @@ class OpenTelemetry(CustomLogger): ) # Create a temporary tracer provider with dynamic headers - temp_provider = TracerProvider(resource=self._get_litellm_resource(self.config)) + temp_provider = TracerProvider( + resource=self._get_litellm_resource(self.config) + ) temp_provider.add_span_processor( self._get_span_processor(dynamic_headers=dynamic_headers) ) @@ -755,7 +784,9 @@ class OpenTelemetry(CustomLogger): metadata = litellm_params.get("metadata") or {} generation_name = metadata.get("generation_name") - raw_span_name = generation_name if generation_name else RAW_REQUEST_SPAN_NAME + raw_span_name = ( + generation_name if generation_name else RAW_REQUEST_SPAN_NAME + ) otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs) raw_span = otel_tracer.start_span( @@ -780,7 +811,9 @@ class OpenTelemetry(CustomLogger): } std_log = kwargs.get("standard_logging_object") - md = getattr(std_log, "metadata", None) or (std_log or {}).get("metadata", {}) + md = getattr(std_log, "metadata", None) or (std_log or {}).get( + "metadata", {} + ) for key in [ "user_api_key_hash", "user_api_key_alias", @@ -802,9 +835,9 @@ class OpenTelemetry(CustomLogger): common_attrs[f"metadata.{key}"] = str(md[key]) # get hidden params - hidden_params = getattr(std_log, "hidden_params", None) or (std_log or {}).get( - "hidden_params", {} - ) + hidden_params = getattr(std_log, "hidden_params", None) or ( + std_log or {} + ).get("hidden_params", {}) if hidden_params: common_attrs["hidden_params"] = safe_dumps(hidden_params) @@ -838,7 +871,9 @@ class OpenTelemetry(CustomLogger): self._record_response_duration_metric(kwargs, end_time, common_attrs) @staticmethod - def _to_timestamp(val: Optional[Union[datetime, float, str]]) -> Optional[float]: + def _to_timestamp( + val: Optional[Union[datetime, float, str]], + ) -> Optional[float]: """Convert datetime/float/string to timestamp.""" if val is None: return None @@ -855,7 +890,9 @@ class OpenTelemetry(CustomLogger): except ValueError: return None - def _record_time_to_first_token_metric(self, kwargs: dict, common_attrs: dict): + def _record_time_to_first_token_metric( + self, kwargs: dict, common_attrs: dict + ): """Record Time to First Token (TTFT) metric for streaming requests.""" optional_params = kwargs.get("optional_params", {}) is_streaming = optional_params.get("stream", False) @@ -868,7 +905,10 @@ class OpenTelemetry(CustomLogger): api_call_start_time = kwargs.get("api_call_start_time", None) completion_start_time = kwargs.get("completion_start_time", None) - if api_call_start_time is not None and completion_start_time is not None: + if ( + api_call_start_time is not None + and completion_start_time is not None + ): # Convert to timestamps if needed (handles datetime, float, and string) api_call_start_ts = self._to_timestamp(api_call_start_time) completion_start_ts = self._to_timestamp(completion_start_time) @@ -876,7 +916,9 @@ class OpenTelemetry(CustomLogger): if api_call_start_ts is None or completion_start_ts is None: return # Skip recording if conversion failed - time_to_first_token_seconds = completion_start_ts - api_call_start_ts + time_to_first_token_seconds = ( + completion_start_ts - api_call_start_ts + ) self._time_to_first_token_histogram.record( time_to_first_token_seconds, attributes=common_attrs ) @@ -946,7 +988,9 @@ class OpenTelemetry(CustomLogger): generation_time_seconds = duration_s if generation_time_seconds > 0: - time_per_output_token_seconds = generation_time_seconds / completion_tokens + time_per_output_token_seconds = ( + generation_time_seconds / completion_tokens + ) self._time_per_output_token_histogram.record( time_per_output_token_seconds, attributes=common_attrs ) @@ -1007,21 +1051,26 @@ class OpenTelemetry(CustomLogger): # See: https://github.com/open-telemetry/opentelemetry-python/pull/4676 # TODO: Refactor to use the proper OTEL Logs API instead of directly creating SDK LogRecords - from opentelemetry._logs import SeverityNumber, get_logger, get_logger_provider + from opentelemetry._logs import ( + SeverityNumber, + get_logger, + ) - try: - from opentelemetry.sdk._logs import LogRecord as SdkLogRecord # type: ignore[attr-defined] # OTEL < 1.39.0 - except ImportError: - from opentelemetry.sdk._logs._internal import LogRecord as SdkLogRecord # type: ignore[attr-defined, no-redef] # OTEL >= 1.39.0 + # MyPy evaluates both branches of try/except imports and can fail when + # newer OTEL stubs remove/relocate symbols. Gate the typing import so + # only the canonical location is type-checked. + if TYPE_CHECKING: + from opentelemetry.sdk._logs._internal import LogRecord as SdkLogRecord + else: + try: + from opentelemetry.sdk._logs import ( + LogRecord as SdkLogRecord, # type: ignore[attr-defined] + ) + except ImportError: + from opentelemetry.sdk._logs._internal import LogRecord as SdkLogRecord otel_logger = get_logger(LITELLM_LOGGER_NAME) - # Get the resource from the logger provider - logger_provider = get_logger_provider() - resource = getattr( - logger_provider, "_resource", None - ) or self._get_litellm_resource(self.config) - parent_ctx = span.get_span_context() provider = (kwargs.get("litellm_params") or {}).get( "custom_llm_provider", "Unknown" @@ -1030,7 +1079,10 @@ class OpenTelemetry(CustomLogger): # per-message events for msg in kwargs.get("messages", []): role = msg.get("role", "user") - attrs = {"event_name": "gen_ai.content.prompt", "gen_ai.system": provider} + attrs = { + "event_name": "gen_ai.content.prompt", + "gen_ai.system": provider, + } if role == "tool" and msg.get("id"): attrs["id"] = msg["id"] if self.message_logging and msg.get("content"): @@ -1044,7 +1096,6 @@ class OpenTelemetry(CustomLogger): severity_number=SeverityNumber.INFO, severity_text="INFO", body=msg.copy(), - resource=resource, attributes=attrs, ) otel_logger.emit(log_record) @@ -1076,7 +1127,6 @@ class OpenTelemetry(CustomLogger): severity_number=SeverityNumber.INFO, severity_text="INFO", body=body, - resource=resource, attributes=attrs, ) otel_logger.emit(log_record) @@ -1146,7 +1196,9 @@ class OpenTelemetry(CustomLogger): value=guardrail_information.get("guardrail_mode"), ) - masked_entity_count = guardrail_information.get("masked_entity_count") + masked_entity_count = guardrail_information.get( + "masked_entity_count" + ) if masked_entity_count is not None: guardrail_span.set_attribute( "masked_entity_count", safe_dumps(masked_entity_count) @@ -1173,8 +1225,9 @@ class OpenTelemetry(CustomLogger): # Decide whether to create a primary span # Always create if no parent span exists (backward compatibility) # OR if USE_OTEL_LITELLM_REQUEST_SPAN is explicitly enabled - should_create_primary_span = parent_otel_span is None or get_secret_bool( - "USE_OTEL_LITELLM_REQUEST_SPAN" + should_create_primary_span = ( + parent_otel_span is None + or get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN") ) if should_create_primary_span: @@ -1200,7 +1253,9 @@ class OpenTelemetry(CustomLogger): if parent_otel_span.is_recording(): parent_otel_span.set_status(Status(StatusCode.ERROR)) self.set_attributes(parent_otel_span, kwargs, response_obj) - self._record_exception_on_span(span=parent_otel_span, kwargs=kwargs) + self._record_exception_on_span( + span=parent_otel_span, kwargs=kwargs + ) # Create span for guardrail information self._create_guardrail_span(kwargs=kwargs, context=_parent_context) @@ -1223,7 +1278,9 @@ class OpenTelemetry(CustomLogger): 2. Sets structured error attributes from StandardLoggingPayloadErrorInformation """ try: - from litellm.integrations._types.open_inference import ErrorAttributes + from litellm.integrations._types.open_inference import ( + ErrorAttributes, + ) # Get the exception object if available exception = kwargs.get("exception") @@ -1233,15 +1290,17 @@ class OpenTelemetry(CustomLogger): span.record_exception(exception) # Get StandardLoggingPayload for structured error information - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" + standard_logging_payload: Optional[StandardLoggingPayload] = ( + kwargs.get("standard_logging_object") ) if standard_logging_payload is None: return # Extract error_information from StandardLoggingPayload - error_information = standard_logging_payload.get("error_information") + error_information = standard_logging_payload.get( + "error_information" + ) if error_information is None: # Fallback to error_str if error_information is not available @@ -1331,7 +1390,9 @@ class OpenTelemetry(CustomLogger): ) pass - def cast_as_primitive_value_type(self, value) -> Union[str, bool, int, float]: + def cast_as_primitive_value_type( + self, value + ) -> Union[str, bool, int, float]: """ Casts the value to a primitive OTEL type if it is not already a primitive type. @@ -1401,8 +1462,8 @@ class OpenTelemetry(CustomLogger): optional_params = kwargs.get("optional_params", {}) litellm_params = kwargs.get("litellm_params", {}) or {} - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" + standard_logging_payload: Optional[StandardLoggingPayload] = ( + kwargs.get("standard_logging_object") ) if standard_logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") @@ -1424,11 +1485,13 @@ class OpenTelemetry(CustomLogger): ) or (standard_logging_payload or {}).get("hidden_params", {}) if hidden_params: self.safe_set_attribute( - span=span, key="hidden_params", value=safe_dumps(hidden_params) + span=span, + key="hidden_params", + value=safe_dumps(hidden_params), ) # Cost breakdown tracking - cost_breakdown: Optional[CostBreakdown] = standard_logging_payload.get( - "cost_breakdown" + cost_breakdown: Optional[CostBreakdown] = ( + standard_logging_payload.get("cost_breakdown") ) if cost_breakdown: for key, value in cost_breakdown.items(): @@ -1504,7 +1567,9 @@ class OpenTelemetry(CustomLogger): # The unique identifier for the completion. if response_obj and response_obj.get("id"): self.safe_set_attribute( - span=span, key="gen_ai.response.id", value=response_obj.get("id") + span=span, + key="gen_ai.response.id", + value=response_obj.get("id"), ) # The model used to generate the response. @@ -1639,7 +1704,9 @@ class OpenTelemetry(CustomLogger): "OpenTelemetry logging error in set_attributes %s", str(e) ) - def _cast_as_primitive_value_type(self, value) -> Union[str, bool, int, float]: + def _cast_as_primitive_value_type( + self, value + ) -> Union[str, bool, int, float]: """ Casts the value to a primitive OTEL type if it is not already a primitive type. @@ -1673,7 +1740,10 @@ class OpenTelemetry(CustomLogger): if isinstance(messages, str): # Handle system_instructions passed as a string return [ - {"role": "system", "parts": [{"type": "text", "content": messages}]} + { + "role": "system", + "parts": [{"type": "text", "content": messages}], + } ] transformed = [] @@ -1714,9 +1784,11 @@ class OpenTelemetry(CustomLogger): message = choice.get("message") or {} finish_reason = choice.get("finish_reason") - transformed_msg = self._transform_messages_to_otel_semantic_conventions( - [message] - )[0] + transformed_msg = ( + self._transform_messages_to_otel_semantic_conventions( + [message] + )[0] + ) if finish_reason: transformed_msg["finish_reason"] = finish_reason @@ -1728,7 +1800,9 @@ class OpenTelemetry(CustomLogger): self.set_attributes(span, kwargs, response_obj) kwargs.get("optional_params", {}) litellm_params = kwargs.get("litellm_params", {}) or {} - custom_llm_provider = litellm_params.get("custom_llm_provider", "Unknown") + custom_llm_provider = litellm_params.get( + "custom_llm_provider", "Unknown" + ) _raw_response = kwargs.get("original_response") _additional_args = kwargs.get("additional_args", {}) or {} @@ -1741,7 +1815,9 @@ class OpenTelemetry(CustomLogger): if complete_input_dict and isinstance(complete_input_dict, dict): for param, val in complete_input_dict.items(): self.safe_set_attribute( - span=span, key=f"llm.{custom_llm_provider}.{param}", value=val + span=span, + key=f"llm.{custom_llm_provider}.{param}", + value=val, ) ############################################# @@ -1773,7 +1849,8 @@ class OpenTelemetry(CustomLogger): ) except Exception as e: verbose_logger.exception( - "OpenTelemetry logging error in set_raw_request_attributes %s", str(e) + "OpenTelemetry logging error in set_raw_request_attributes %s", + str(e), ) def _to_ns(self, dt): @@ -1813,7 +1890,9 @@ class OpenTelemetry(CustomLogger): ) litellm_params = kwargs.get("litellm_params", {}) or {} - proxy_server_request = litellm_params.get("proxy_server_request", {}) or {} + proxy_server_request = ( + litellm_params.get("proxy_server_request", {}) or {} + ) headers = proxy_server_request.get("headers", {}) or {} traceparent = headers.get("traceparent", None) _metadata = litellm_params.get("metadata", {}) or {} @@ -1832,7 +1911,10 @@ class OpenTelemetry(CustomLogger): "OpenTelemetry: Using traceparent header for context propagation" ) carrier = {"traceparent": traceparent} - return TraceContextTextMapPropagator().extract(carrier=carrier), None + return ( + TraceContextTextMapPropagator().extract(carrier=carrier), + None, + ) # Priority 3: Active span from global context (auto-detection) try: @@ -1960,10 +2042,14 @@ class OpenTelemetry(CustomLogger): self.OTEL_HEADERS, ) - _split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS) + _split_otel_headers = OpenTelemetry._get_headers_dictionary( + self.OTEL_HEADERS + ) # Normalize endpoint for logs - ensure it points to /v1/logs instead of /v1/traces - normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "logs") + normalized_endpoint = self._normalize_otel_endpoint( + self.OTEL_ENDPOINT, "logs" + ) verbose_logger.debug( "OpenTelemetry: Log endpoint normalized from %s to %s", @@ -2051,14 +2137,18 @@ class OpenTelemetry(CustomLogger): self.OTEL_HEADERS, ) - _split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS) + _split_otel_headers = OpenTelemetry._get_headers_dictionary( + self.OTEL_HEADERS + ) normalized_endpoint = self._normalize_otel_endpoint( self.OTEL_ENDPOINT, "metrics" ) if self.OTEL_EXPORTER == "console": exporter = ConsoleMetricExporter() - return PeriodicExportingMetricReader(exporter, export_interval_millis=5000) + return PeriodicExportingMetricReader( + exporter, export_interval_millis=5000 + ) elif ( self.OTEL_EXPORTER == "otlp_http" @@ -2074,7 +2164,9 @@ class OpenTelemetry(CustomLogger): headers=_split_otel_headers, preferred_temporality={Histogram: AggregationTemporality.DELTA}, ) - return PeriodicExportingMetricReader(exporter, export_interval_millis=5000) + return PeriodicExportingMetricReader( + exporter, export_interval_millis=5000 + ) elif self.OTEL_EXPORTER == "otlp_grpc" or self.OTEL_EXPORTER == "grpc": try: @@ -2092,7 +2184,9 @@ class OpenTelemetry(CustomLogger): headers=_split_otel_headers, preferred_temporality={Histogram: AggregationTemporality.DELTA}, ) - return PeriodicExportingMetricReader(exporter, export_interval_millis=5000) + return PeriodicExportingMetricReader( + exporter, export_interval_millis=5000 + ) else: verbose_logger.warning( @@ -2100,7 +2194,9 @@ class OpenTelemetry(CustomLogger): self.OTEL_EXPORTER, ) exporter = ConsoleMetricExporter() - return PeriodicExportingMetricReader(exporter, export_interval_millis=5000) + return PeriodicExportingMetricReader( + exporter, export_interval_millis=5000 + ) def _normalize_otel_endpoint( self, endpoint: Optional[str], signal_type: str @@ -2171,7 +2267,9 @@ class OpenTelemetry(CustomLogger): return endpoint @staticmethod - def _get_headers_dictionary(headers: Optional[Union[str, dict]]) -> Dict[str, str]: + def _get_headers_dictionary( + headers: Optional[Union[str, dict]], + ) -> Dict[str, str]: """ Convert a string or dictionary of headers into a dictionary of headers. """ diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 9b86f4ca2f0..e622a317454 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -8,40 +8,187 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True ``` """ +import json import os +from importlib.resources import files import httpx +from litellm import verbose_logger +from litellm.constants import ( + MODEL_COST_MAP_MAX_SHRINK_RATIO, + MODEL_COST_MAP_MIN_MODEL_COUNT, +) + + +class GetModelCostMap: + """ + Handles fetching, validating, and loading the model cost map. + + Only the backup model *count* is cached (a single int). The full + backup dict is never held in memory — it is only parsed when it + needs to be *returned* as a fallback. + """ + + _backup_model_count: int = -1 # -1 = not yet loaded + + @staticmethod + def load_local_model_cost_map() -> dict: + """Load the local backup model cost map bundled with the package.""" + content = json.loads( + files("litellm") + .joinpath("model_prices_and_context_window_backup.json") + .read_text(encoding="utf-8") + ) + return content + + @classmethod + def _get_backup_model_count(cls) -> int: + """Return the number of models in the local backup (cached int).""" + if cls._backup_model_count < 0: + backup = cls.load_local_model_cost_map() + cls._backup_model_count = len(backup) + return cls._backup_model_count + + @staticmethod + def _check_is_valid_dict(fetched_map: dict) -> bool: + """Check 1: fetched map is a non-empty dict.""" + if not isinstance(fetched_map, dict): + verbose_logger.warning( + "LiteLLM: Fetched model cost map is not a dict (type=%s). " + "Falling back to local backup.", + type(fetched_map).__name__, + ) + return False + + if len(fetched_map) == 0: + verbose_logger.warning( + "LiteLLM: Fetched model cost map is empty. " + "Falling back to local backup.", + ) + return False + + return True + + @classmethod + def _check_model_count_not_reduced( + cls, + fetched_map: dict, + backup_model_count: int, + min_model_count: int = MODEL_COST_MAP_MIN_MODEL_COUNT, + max_shrink_ratio: float = MODEL_COST_MAP_MAX_SHRINK_RATIO, + ) -> bool: + """Check 2: model count has not reduced significantly vs backup.""" + fetched_count = len(fetched_map) + + if fetched_count < min_model_count: + verbose_logger.warning( + "LiteLLM: Fetched model cost map has only %d models (minimum=%d). " + "This may indicate a corrupted upstream file. " + "Falling back to local backup.", + fetched_count, + min_model_count, + ) + return False + + if backup_model_count > 0 and fetched_count < backup_model_count * max_shrink_ratio: + verbose_logger.warning( + "LiteLLM: Fetched model cost map shrank significantly " + "(fetched=%d, backup=%d, threshold=%.0f%%). " + "This may indicate a corrupted upstream file. " + "Falling back to local backup.", + fetched_count, + backup_model_count, + max_shrink_ratio * 100, + ) + return False + + return True + + @classmethod + def validate_model_cost_map( + cls, + fetched_map: dict, + backup_model_count: int, + min_model_count: int = MODEL_COST_MAP_MIN_MODEL_COUNT, + max_shrink_ratio: float = MODEL_COST_MAP_MAX_SHRINK_RATIO, + ) -> bool: + """ + Validate the integrity of a fetched model cost map. + + Runs each check in order and returns False on the first failure. + + Checks: + 1. ``_check_is_valid_dict`` -- fetched map is a non-empty dict. + 2. ``_check_model_count_not_reduced`` -- model count meets minimum + and has not shrunk >``max_shrink_ratio`` vs backup. + + Returns True if all checks pass, False otherwise. + """ + if not cls._check_is_valid_dict(fetched_map): + return False + + if not cls._check_model_count_not_reduced( + fetched_map=fetched_map, + backup_model_count=backup_model_count, + min_model_count=min_model_count, + max_shrink_ratio=max_shrink_ratio, + ): + return False + + return True + + @staticmethod + def fetch_remote_model_cost_map(url: str, timeout: int = 5) -> dict: + """ + Fetch the model cost map from a remote URL. + + Returns the parsed JSON dict. Raises on network/parse errors + (caller is expected to handle). + """ + response = httpx.get(url, timeout=timeout) + response.raise_for_status() + return response.json() + def get_model_cost_map(url: str) -> dict: - if ( - os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", False) - or os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", False) == "True" - ): - from importlib.resources import files - import json + """ + Public entry point — returns the model cost map dict. - content = json.loads( - files("litellm") - .joinpath("model_prices_and_context_window_backup.json") - .read_text(encoding="utf-8") - ) - return content + 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. + 2. Otherwise fetches from ``url``, validates integrity, and falls back + to the local backup on any failure. + + Only the backup model count is cached (a single int) for validation. + The full backup dict is only parsed when it must be *returned* as a + fallback — it is never held in memory long-term. + """ + # 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": + return GetModelCostMap.load_local_model_cost_map() try: - response = httpx.get( - url, timeout=5 - ) # set a 5 second timeout for the get request - response.raise_for_status() # Raise an exception if the request is unsuccessful - content = response.json() - return content - except Exception: - from importlib.resources import files - import json - - content = json.loads( - files("litellm") - .joinpath("model_prices_and_context_window_backup.json") - .read_text(encoding="utf-8") + content = GetModelCostMap.fetch_remote_model_cost_map(url) + except Exception as e: + verbose_logger.warning( + "LiteLLM: Failed to fetch remote model cost map from %s: %s. " + "Falling back to local backup.", + url, + str(e), ) - return content + return GetModelCostMap.load_local_model_cost_map() + + # Validate using cached count (cheap int comparison, no file I/O) + if not GetModelCostMap.validate_model_cost_map( + fetched_map=content, + backup_model_count=GetModelCostMap._get_backup_model_count(), + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. " + "Using local backup instead. url=%s", + url, + ) + return GetModelCostMap.load_local_model_cost_map() + + return content diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7ae62718af0..82a7af64f97 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3764,7 +3764,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 from litellm.integrations.opentelemetry import OpenTelemetry for callback in _in_memory_loggers: - if isinstance(callback, OpenTelemetry): + if type(callback) is OpenTelemetry: return callback # type: ignore otel_logger = OpenTelemetry( **_get_custom_logger_settings_from_proxy_server( diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index b1c2d0a52f5..cdddee4e54e 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1272,3 +1272,59 @@ def parse_tool_call_arguments( ) raise ValueError(error_message) from e + + +def split_concatenated_json_objects(raw: str) -> List[Dict[str, Any]]: + """ + Split a string that contains one or more concatenated JSON objects into + a list of parsed dicts. + + LLM providers (notably Bedrock Claude Sonnet 4.5) sometimes return + multiple tool-call argument objects concatenated in a single + ``arguments`` string, e.g.:: + + '{"command":["curl",...]}{"command":["curl",...]}{"command":["curl",...]}' + + ``json.loads()`` fails on this with ``JSONDecodeError: Extra data``. + This helper uses ``json.JSONDecoder.raw_decode()`` to walk the string + and extract each JSON object individually. + + Returns + ------- + list[dict] + A list of parsed dicts – one per JSON object found. If *raw* is + empty or whitespace-only, an empty list is returned. + + Raises + ------ + json.JSONDecodeError + If the string contains text that cannot be parsed as JSON at all. + """ + import json + + raw = raw.strip() + if not raw: + return [] + + decoder = json.JSONDecoder() + results: List[Dict[str, Any]] = [] + idx = 0 + length = len(raw) + + while idx < length: + # Skip whitespace between objects + while idx < length and raw[idx] in " \t\n\r": + idx += 1 + if idx >= length: + break + + obj, end_idx = decoder.raw_decode(raw, idx) + if isinstance(obj, dict): + results.append(obj) + else: + # Non-dict JSON value – wrap in empty dict (Bedrock requires + # toolUse.input to be an object). + results.append({}) + idx = end_idx + + return results diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index f9ecd78ff1c..c907ed32b95 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3287,25 +3287,68 @@ def _convert_to_bedrock_tool_call_invoke( - extract name - extract id """ + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + split_concatenated_json_objects, + ) try: _parts_list: List[BedrockContentBlock] = [] for tool in tool_calls: if "function" in tool: - id = tool["id"] + tool_id = tool["id"] name = tool["function"].get("name", "") arguments = tool["function"].get("arguments", "") - arguments_dict = json.loads(arguments) if arguments else {} - # Ensure arguments_dict is always a dict (Bedrock requires toolUse.input to be an object) - # When some providers return arguments: '""' (JSON-encoded empty string), json.loads returns "" - if not isinstance(arguments_dict, dict): - arguments_dict = {} + if not arguments or not arguments.strip(): arguments_dict = {} else: - arguments_dict = json.loads(arguments) + try: + arguments_dict = 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): + arguments_dict = {} + except json.JSONDecodeError: + # The model may return multiple JSON objects + # concatenated in a single arguments string, e.g. + # '{"cmd":"a"}{"cmd":"b"}{"cmd":"c"}' + # Split them and emit one toolUse block per object. + # Fixes: https://github.com/BerriAI/litellm/issues/20543 + parsed_objects = split_concatenated_json_objects( + arguments + ) + if parsed_objects: + # First object keeps the original tool id. + for obj_idx, obj in enumerate(parsed_objects): + block_id = ( + tool_id + if obj_idx == 0 + else f"{tool_id}_{obj_idx}" + ) + bedrock_tool = BedrockToolUseBlock( + input=obj, name=name, toolUseId=block_id + ) + _parts_list.append( + BedrockContentBlock(toolUse=bedrock_tool) + ) + # cache_control applies to the whole original + # tool call; attach after the last split block. + if tool.get("cache_control", None) is not None: + _parts_list.append( + BedrockContentBlock( + cachePoint=CachePointBlock( + type="default" + ) + ) + ) + continue + # Fallback: no objects extracted — use empty dict. + arguments_dict = {} + bedrock_tool = BedrockToolUseBlock( - input=arguments_dict, name=name, toolUseId=id + input=arguments_dict, name=name, toolUseId=tool_id ) bedrock_content_block = BedrockContentBlock(toolUse=bedrock_tool) _parts_list.append(bedrock_content_block) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index aa763dc9899..5d6d1fbc1c5 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -140,9 +140,14 @@ def should_redact_message_logging(model_call_details: dict) -> bool: metadata_field = get_metadata_variable_name_from_kwargs(litellm_params) metadata = litellm_params.get(metadata_field, {}) - + if not isinstance(metadata, dict): + # Fall back: litellm_metadata was None, try metadata + metadata = litellm_params.get("metadata", {}) + if not isinstance(metadata, dict): + metadata = {} + # Get headers from the metadata - request_headers = metadata.get("headers", {}) if isinstance(metadata, dict) else {} + request_headers = metadata.get("headers", {}) # Check for headers that explicitly control redaction if request_headers and bool( diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index e85c0d0d017..f51adf96102 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -58,6 +58,9 @@ from litellm.types.utils import ( from ...base import BaseLLM from ..common_utils import AnthropicError, process_anthropic_headers +from litellm.anthropic_beta_headers_manager import ( + update_headers_with_filtered_beta, +) from .transformation import AnthropicConfig if TYPE_CHECKING: @@ -333,6 +336,10 @@ class AnthropicChatCompletion(BaseLLM): litellm_params=litellm_params, ) + headers = update_headers_with_filtered_beta( + headers=headers, provider=custom_llm_provider + ) + config = ProviderConfigManager.get_provider_chat_config( model=model, provider=LlmProviders(custom_llm_provider), diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 82aa7390188..acd08f0a569 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -834,6 +834,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "sonnet-4-5", "opus-4.1", "opus-4-1", + "opus-4.5", + "opus-4-5", + "opus-4.6", + "opus-4-6", } ): _output_format = ( @@ -931,6 +935,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): Translate system message to anthropic format. Removes system message from the original list and returns a new list of anthropic system message content. + Filters out system messages containing x-anthropic-billing-header metadata. """ system_prompt_indices = [] anthropic_system_message_list: List[AnthropicSystemMessageContent] = [] @@ -942,6 +947,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Skip empty text blocks - Anthropic API raises errors for empty text if not system_message_block["content"]: continue + # Skip system messages containing x-anthropic-billing-header metadata + if system_message_block["content"].startswith("x-anthropic-billing-header:"): + continue anthropic_system_message_content = AnthropicSystemMessageContent( type="text", text=system_message_block["content"], @@ -960,6 +968,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): text_value = _content.get("text") if _content.get("type") == "text" and not text_value: continue + # Skip system messages containing x-anthropic-billing-header metadata + if _content.get("type") == "text" and text_value and text_value.startswith("x-anthropic-billing-header:"): + continue anthropic_system_message_content = ( AnthropicSystemMessageContent( type=_content.get("type"), diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a17eba75b3b..c6caaddf98b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -30,6 +30,58 @@ ANTHROPIC_ADAPTER = AnthropicAdapter() class LiteLLMMessagesToCompletionTransformationHandler: + @staticmethod + def _route_openai_thinking_to_responses_api_if_needed( + completion_kwargs: Dict[str, Any], + *, + thinking: Optional[Dict[str, Any]], + ) -> None: + """ + When users call `litellm.anthropic.messages.*` with a non-Anthropic model and + `thinking={"type": "enabled", ...}`, LiteLLM converts this into OpenAI + `reasoning_effort`. + + For OpenAI models, Chat Completions typically does not return reasoning text + (only token accounting). To return a thinking-like content block in the + Anthropic response format, we route the request through OpenAI's Responses API + and request a reasoning summary. + """ + custom_llm_provider = completion_kwargs.get("custom_llm_provider") + if custom_llm_provider is None: + try: + _, inferred_provider, _, _ = litellm.utils.get_llm_provider( + model=cast(str, completion_kwargs.get("model")) + ) + custom_llm_provider = inferred_provider + except Exception: + custom_llm_provider = None + + if custom_llm_provider != "openai": + return + + if not isinstance(thinking, dict) or thinking.get("type") != "enabled": + return + + model = completion_kwargs.get("model") + if isinstance(model, str) and model and not model.startswith("responses/"): + # Prefix model with "responses/" to route to OpenAI Responses API + completion_kwargs["model"] = f"responses/{model}" + + reasoning_effort = completion_kwargs.get("reasoning_effort") + if isinstance(reasoning_effort, str) and reasoning_effort: + completion_kwargs["reasoning_effort"] = { + "effort": reasoning_effort, + "summary": "detailed", + } + elif isinstance(reasoning_effort, dict): + if ( + "summary" not in reasoning_effort + and "generate_summary" not in reasoning_effort + ): + updated_reasoning_effort = dict(reasoning_effort) + updated_reasoning_effort["summary"] = "detailed" + completion_kwargs["reasoning_effort"] = updated_reasoning_effort + @staticmethod def _prepare_completion_kwargs( *, @@ -123,6 +175,11 @@ class LiteLLMMessagesToCompletionTransformationHandler: ): completion_kwargs[key] = value + LiteLLMMessagesToCompletionTransformationHandler._route_openai_thinking_to_responses_api_if_needed( + completion_kwargs, + thinking=thinking, + ) + return completion_kwargs, tool_name_mapping @staticmethod diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 043a70f3c67..8f2f3bf3545 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -2,9 +2,6 @@ from typing import Any, AsyncIterator, Dict, List, Optional, Tuple import httpx -from litellm.anthropic_beta_headers_manager import ( - update_headers_with_filtered_beta, -) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import verbose_logger from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -52,6 +49,40 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # TODO: Add Anthropic `metadata` support # "metadata", ] + + @staticmethod + def _filter_billing_headers_from_system(system_param): + """ + Filter out x-anthropic-billing-header metadata from system parameter. + + Args: + system_param: Can be a string or a list of system message content blocks + + Returns: + Filtered system parameter (string or list), or None if all content was filtered + """ + if isinstance(system_param, str): + # If it's a string and starts with billing header, filter it out + if system_param.startswith("x-anthropic-billing-header:"): + return None + return system_param + elif isinstance(system_param, list): + # Filter list of system content blocks + filtered_list = [] + for content_block in system_param: + if isinstance(content_block, dict): + text = content_block.get("text", "") + content_type = content_block.get("type", "") + # Skip text blocks that start with billing header + if content_type == "text" and text.startswith("x-anthropic-billing-header:"): + continue + filtered_list.append(content_block) + else: + # Keep non-dict items as-is + filtered_list.append(content_block) + return filtered_list if len(filtered_list) > 0 else None + else: + return system_param def get_complete_url( self, @@ -96,11 +127,6 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params=optional_params, ) - headers = update_headers_with_filtered_beta( - headers=headers, - provider="anthropic", - ) - return headers, api_base def transform_anthropic_messages_request( @@ -123,6 +149,17 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): message="max_tokens is required for Anthropic /v1/messages API", status_code=400, ) + + # Filter out x-anthropic-billing-header from system messages + system_param = anthropic_messages_optional_request_params.get("system") + if system_param is not None: + filtered_system = self._filter_billing_headers_from_system(system_param) + if filtered_system is not None and len(filtered_system) > 0: + anthropic_messages_optional_request_params["system"] = filtered_system + else: + # Remove system parameter if all content was filtered out + anthropic_messages_optional_request_params.pop("system", None) + ####### get required params for all anthropic messages requests ###### verbose_logger.debug(f"TRANSFORMATION DEBUG - Messages: {messages}") anthropic_messages_request: AnthropicMessagesRequest = AnthropicMessagesRequest( diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 0ae6fad7300..18dad503a59 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -105,6 +105,7 @@ class AzureOpenAIConfig(BaseConfig): "modalities", "audio", "web_search_options", + "prompt_cache_key", ] def _is_response_format_supported_model(self, model: str) -> bool: diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index f86ec7082f2..a4dc88f9c68 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -3,9 +3,6 @@ Azure Anthropic messages transformation config - extends AnthropicMessagesConfig """ from typing import TYPE_CHECKING, Any, List, Optional, Tuple -from litellm.anthropic_beta_headers_manager import ( - update_headers_with_filtered_beta, -) from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -65,18 +62,11 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): if "content-type" not in headers: headers["content-type"] = "application/json" - # Update headers with anthropic beta features (context management, tool search, etc.) headers = self._update_headers_with_anthropic_beta( headers=headers, optional_params=optional_params, ) - # Filter out unsupported beta headers for Azure AI - headers = update_headers_with_filtered_beta( - headers=headers, - provider="azure_ai", - ) - return headers, api_base def get_complete_url( diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index 753bc9c08eb..c5510db68b1 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -2,10 +2,6 @@ Azure Anthropic transformation config - extends AnthropicConfig with Azure authentication """ from typing import TYPE_CHECKING, Dict, List, Optional, Union - -from litellm.anthropic_beta_headers_manager import ( - update_headers_with_filtered_beta, -) from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.types.llms.openai import AllMessageValues @@ -90,11 +86,6 @@ class AzureAnthropicConfig(AnthropicConfig): if "anthropic-version" not in headers: headers["anthropic-version"] = "2023-06-01" - # Filter out unsupported beta headers for Azure AI - headers = update_headers_with_filtered_beta( - headers=headers, - provider="azure_ai", - ) return headers diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 04d2b3a2769..585efd3307d 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -11,12 +11,14 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( _audio_or_image_in_message_content, convert_content_list_to_str, ) +from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.xai.chat.transformation import XAIChatConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ModelResponse, ProviderField from litellm.utils import _add_path_to_api_base, supports_tool_choice @@ -64,12 +66,21 @@ class AzureAIStudioConfig(OpenAIConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - if api_base and self._should_use_api_key_header(api_base): - headers["api-key"] = api_key + if api_key: + if api_base and self._should_use_api_key_header(api_base): + headers["api-key"] = api_key + else: + headers["Authorization"] = f"Bearer {api_key}" else: - headers["Authorization"] = f"Bearer {api_key}" + # No api_key provided — fall back to Azure AD token-based auth + litellm_params_obj = GenericLiteLLMParams( + **(litellm_params if isinstance(litellm_params, dict) else {}) + ) + headers = BaseAzureLLM._base_validate_azure_environment( + headers=headers, litellm_params=litellm_params_obj + ) - headers["Content-Type"] = "application/json" # tell Azure AI Studio to expect JSON + headers["Content-Type"] = "application/json" return headers diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1de1c40c438..304c707fa0b 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -211,25 +211,13 @@ class BaseAWSLLM: aws_external_id=aws_external_id, ) elif aws_role_name is not None: - # Check if we're in IRSA and trying to assume the same role we already have - current_role_arn = os.getenv("AWS_ROLE_ARN") - web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") - - # In IRSA environments, we should skip role assumption if we're already running as the target role - # This is true when: - # 1. We have AWS_ROLE_ARN set (current role) - # 2. We have AWS_WEB_IDENTITY_TOKEN_FILE set (IRSA environment) - # 3. The current role matches the requested role - if ( - current_role_arn - and web_identity_token_file - and current_role_arn == aws_role_name - ): + # Check if we're already running as the target role and can skip assumption + # This handles IRSA (EKS), ECS task roles, and EC2 instance profiles + if self._is_already_running_as_role(aws_role_name, ssl_verify=ssl_verify): verbose_logger.debug( - "Using IRSA same-role optimization: calling _auth_with_env_vars" + "Already running as target role %s, using ambient credentials", + aws_role_name, ) - # We're already running as this role via IRSA, no need to assume it again - # Use the default boto3 credentials (which will use the IRSA credentials) credentials, _cache_ttl = self._auth_with_env_vars() else: verbose_logger.debug( @@ -553,6 +541,107 @@ class BaseAWSLLM: aws_region_name = "us-west-2" return aws_region_name + @staticmethod + def _parse_arn_account_and_role_name( + arn: str, + ) -> Optional[Tuple[str, str, str]]: + """ + Parse an ARN and return (partition, account_id, role_name). + + Handles: + - arn:aws:iam::123456789012:role/MyRole + - arn:aws:iam::123456789012:role/path/to/MyRole + - arn:aws:sts::123456789012:assumed-role/MyRole/session-name + + Returns None if the ARN cannot be parsed. + """ + # ARN format: arn:PARTITION:SERVICE:REGION:ACCOUNT:RESOURCE + parts = arn.split(":") + if len(parts) < 6 or parts[0] != "arn": + return None + + partition = parts[1] # e.g. "aws", "aws-cn", "aws-us-gov" + account_id = parts[4] + resource = ":".join(parts[5:]) # rejoin in case resource contains colons + + if resource.startswith("role/"): + # arn:aws:iam::ACCOUNT:role/[path/]ROLE_NAME + role_name = resource.split("/")[-1] + elif resource.startswith("assumed-role/"): + # arn:aws:sts::ACCOUNT:assumed-role/ROLE_NAME/SESSION + role_parts = resource.split("/") + if len(role_parts) >= 2: + role_name = role_parts[1] + else: + return None + else: + return None + + return partition, account_id, role_name + + def _is_already_running_as_role( + self, + aws_role_name: str, + ssl_verify: Optional[Union[bool, str]] = None, + ) -> bool: + """ + Check if the current environment is already running as the target IAM role. + + This handles multiple AWS environments: + - IRSA (EKS): AWS_ROLE_ARN + AWS_WEB_IDENTITY_TOKEN_FILE are set + - ECS task roles: Uses sts:GetCallerIdentity to check current role ARN + - EC2 instance profiles: Uses sts:GetCallerIdentity to check current role ARN + + Compares partition, account ID, and role name to avoid cross-account + false matches. + + Returns True if the current identity matches the target role, meaning + we can skip sts:AssumeRole and use ambient credentials directly. + """ + target_parsed = self._parse_arn_account_and_role_name(aws_role_name) + if target_parsed is None: + return False + + target_partition, target_account, target_role = target_parsed + + # Fast path: IRSA environment check (no API call needed) + current_role_arn = os.getenv("AWS_ROLE_ARN") + web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") + if current_role_arn and web_identity_token_file: + return current_role_arn == aws_role_name + + # For ECS/EC2: call sts:GetCallerIdentity to check if already running as the role + try: + import boto3 + + with tracer.trace("boto3.client(sts).get_caller_identity"): + sts_client = boto3.client( + "sts", verify=self._get_ssl_verify(ssl_verify) + ) + identity = sts_client.get_caller_identity() + caller_arn = identity.get("Arn", "") + + caller_parsed = self._parse_arn_account_and_role_name(caller_arn) + if caller_parsed is not None: + caller_partition, caller_account, caller_role = caller_parsed + if ( + caller_partition == target_partition + and caller_account == target_account + and caller_role == target_role + ): + verbose_logger.debug( + "Current identity already matches target role: %s", + aws_role_name, + ) + return True + + except Exception as e: + verbose_logger.debug( + "Could not determine current role identity: %s", str(e) + ) + + return False + @tracer.wrap() def _auth_with_web_identity_token( self, @@ -867,7 +956,35 @@ class BaseAWSLLM: if aws_external_id is not None: assume_role_params["ExternalId"] = aws_external_id - sts_response = sts_client.assume_role(**assume_role_params) + try: + sts_response = sts_client.assume_role(**assume_role_params) + except Exception as e: + error_str = str(e) + if "AccessDenied" in error_str: + # Only fall back to ambient credentials if we can positively + # confirm the caller is already the target role (same account, + # partition, and role name). This avoids silently using the + # wrong identity when there is a genuine trust-policy or + # permission misconfiguration. + if self._is_already_running_as_role( + aws_role_name, ssl_verify=ssl_verify + ): + verbose_logger.warning( + "AssumeRole failed for %s (%s). " + "Caller is already running as this role; " + "falling back to ambient credentials.", + aws_role_name, + error_str, + ) + return self._auth_with_env_vars() + # Genuine permission error — re-raise + verbose_logger.error( + "AssumeRole AccessDenied for %s and caller is NOT " + "the same role. Re-raising. Error: %s", + aws_role_name, + error_str, + ) + raise # Extract the credentials from the response and convert to Session Credentials sts_credentials = sts_response["Credentials"] diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index d5bd054118d..25af852e09c 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -13,7 +13,9 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper - +from litellm.anthropic_beta_headers_manager import ( + update_headers_with_filtered_beta, + ) from ..base_aws_llm import BaseAWSLLM, Credentials from ..common_utils import BedrockError from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -337,7 +339,11 @@ class BedrockConverseLLM(BaseAWSLLM): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - + + # Filter beta headers in HTTP headers before making the request + headers = update_headers_with_filtered_beta( + headers=headers, provider="bedrock_converse" + ) ### ROUTING (ASYNC, STREAMING, SYNC) if acompletion: if isinstance(client, HTTPHandler): diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 7fc51263ebb..efa755d515e 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -11,9 +11,6 @@ import httpx import litellm from litellm._logging import verbose_logger -from litellm.anthropic_beta_headers_manager import ( - filter_and_transform_beta_headers, -) from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import ( filter_exceptions_from_params, @@ -1132,24 +1129,9 @@ class AmazonConverseConfig(BaseConfig): # Set anthropic_beta in additional_request_params if we have any beta features # ONLY apply to Anthropic/Claude models - other models (e.g., Qwen, Llama) don't support this field - # and will error with "unknown variant anthropic_beta" if included base_model = BedrockModelInfo.get_base_model(model) if anthropic_beta_list and base_model.startswith("anthropic"): - # Remove duplicates while preserving order - unique_betas = [] - seen = set() - for beta in anthropic_beta_list: - if beta not in seen: - unique_betas.append(beta) - seen.add(beta) - - filtered_betas = filter_and_transform_beta_headers( - beta_headers=unique_betas, - provider="bedrock_converse", - ) - - if filtered_betas: - additional_request_params["anthropic_beta"] = filtered_betas + additional_request_params["anthropic_beta"] = anthropic_beta_list return bedrock_tools, anthropic_beta_list diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 31119c73d72..dfab81123fd 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -2,7 +2,6 @@ from typing import TYPE_CHECKING, Any, List, Optional import httpx -from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, @@ -136,13 +135,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Filter out beta headers that Bedrock Invoke doesn't support # Uses centralized configuration from anthropic_beta_headers_config.json beta_list = list(beta_set) - filtered_beta_list = filter_and_transform_beta_headers( - beta_headers=beta_list, - provider="bedrock", - ) - - if filtered_beta_list: - _anthropic_request["anthropic_beta"] = filtered_beta_list + _anthropic_request["anthropic_beta"] = beta_list return _anthropic_request 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 19fe7d8c140..2c041adba7a 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -12,9 +12,6 @@ from typing import ( import httpx -from litellm.anthropic_beta_headers_manager import ( - filter_and_transform_beta_headers, -) from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, @@ -253,71 +250,6 @@ class AmazonAnthropicClaudeMessagesConfig( return any(pattern in model_lower for pattern in supported_patterns) - def _filter_unsupported_beta_headers_for_bedrock( - self, model: str, beta_set: set - ) -> None: - """ - Remove beta headers that are not supported on Bedrock for the given model. - - Extended thinking beta headers are only supported on specific Claude 4+ models. - Advanced tool use headers are not supported on Bedrock Invoke API, but need to be - translated to Bedrock-specific headers for models that support tool search - (Claude Opus 4.5, Sonnet 4.5). - This prevents 400 "invalid beta flag" errors on Bedrock. - - Note: Bedrock Invoke API fails with a 400 error when unsupported beta headers - are sent, returning: {"message":"invalid beta flag"} - - Translation for models supporting tool search (Opus 4.5, Sonnet 4.5): - - advanced-tool-use-2025-11-20 -> tool-search-tool-2025-10-19 + tool-examples-2025-10-29 - - Args: - model: The model name - beta_set: The set of beta headers to filter in-place - """ - # 1. Handle header transformations BEFORE filtering - # (advanced-tool-use -> tool-search-tool) - # This must happen before filtering because advanced-tool-use is in the unsupported list - has_advanced_tool_use = "advanced-tool-use-2025-11-20" in beta_set - if has_advanced_tool_use and self._supports_tool_search_on_bedrock(model): - beta_set.discard("advanced-tool-use-2025-11-20") - beta_set.add("tool-search-tool-2025-10-19") - beta_set.add("tool-examples-2025-10-29") - - # 2. Apply provider-level filtering using centralized JSON config - beta_list = list(beta_set) - filtered_list = filter_and_transform_beta_headers( - beta_headers=beta_list, - provider="bedrock", - ) - - # Update the set with filtered headers - beta_set.clear() - beta_set.update(filtered_list) - - # 2.1. Handle model-specific exceptions: structured-outputs is only supported on Opus 4.6 - # Re-add structured-outputs if it was in the original set and model is Opus 4.6 - model_lower = model.lower() - is_opus_4_6 = any(pattern in model_lower for pattern in ["opus-4.6", "opus_4.6", "opus-4-6", "opus_4_6"]) - if is_opus_4_6 and "structured-outputs-2025-11-13" in beta_list: - beta_set.add("structured-outputs-2025-11-13") - - # 3. Filter out extended thinking headers for models that don't support them - extended_thinking_patterns = [ - "extended-thinking", - "interleaved-thinking", - ] - if not self._supports_extended_thinking_on_bedrock(model): - beta_headers_to_remove = set() - for beta in beta_set: - for pattern in extended_thinking_patterns: - if pattern in beta.lower(): - beta_headers_to_remove.add(beta) - break - - for beta in beta_headers_to_remove: - beta_set.discard(beta) - def _get_tool_search_beta_header_for_bedrock( self, model: str, @@ -483,12 +415,11 @@ class AmazonAnthropicClaudeMessagesConfig( beta_set=beta_set, ) - # Filter out unsupported beta headers for Bedrock (e.g., advanced-tool-use, extended-thinking on non-Opus/Sonnet 4 models) - self._filter_unsupported_beta_headers_for_bedrock( - model=model, - beta_set=beta_set, - ) - + # --- Custom logic: if tool-search-tool-2025-10-19 is present, add tool-examples-2025-10-29 --- + if "tool-search-tool-2025-10-19" in beta_set: + beta_set.add("tool-examples-2025-10-29") + # ------------------------------------------------------------------------------ + if beta_set: anthropic_messages_request["anthropic_beta"] = list(beta_set) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index a7b83d8c802..1b03ec47643 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -1,6 +1,7 @@ import asyncio import contextlib import os +import ssl import typing import urllib.request from typing import Callable, Dict, Optional, Union @@ -139,8 +140,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport): Credit to: https://github.com/karpetrosyan/httpx-aiohttp for this implementation """ - def __init__(self, client: Union[ClientSession, Callable[[], ClientSession]]): + def __init__( + self, + client: Union[ClientSession, Callable[[], ClientSession]], + ssl_verify: Optional[Union[bool, ssl.SSLContext]] = None, + ): self.client = client + self._ssl_verify = ssl_verify # Store for per-request SSL override super().__init__(client=client) # Store the client factory for recreating sessions when needed if callable(client): @@ -214,6 +220,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): timeout: dict, proxy: Optional[str], sni_hostname: Optional[str], + ssl_verify: Optional[Union[bool, ssl.SSLContext]] = None, ) -> ClientResponse: """ Helper function to make an aiohttp request with the given parameters. @@ -224,6 +231,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): timeout: Timeout settings dict with 'connect', 'read', 'pool' keys proxy: Optional proxy URL sni_hostname: Optional SNI hostname for SSL + ssl_verify: Optional SSL verification setting (False to disable, SSLContext for custom) Returns: ClientResponse from aiohttp @@ -237,6 +245,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport): data = request.stream # type: ignore request.headers.pop("transfer-encoding", None) # handled by aiohttp + # Only pass ssl kwarg when explicitly configured, to avoid + # overriding the session/connector defaults with None (which is + # not a valid value for aiohttp's ssl parameter). + ssl_kwargs: Dict[str, Union[bool, ssl.SSLContext]] = {} + if ssl_verify is not None: + ssl_kwargs["ssl"] = ssl_verify + response = await client_session.request( method=request.method, url=YarlURL(str(request.url), encoded=True), @@ -251,6 +266,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): ), proxy=proxy, server_hostname=sni_hostname, + **ssl_kwargs, ).__aenter__() return response @@ -268,6 +284,9 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Resolve proxy settings from environment variables proxy = await self._get_proxy_settings(request) + # Use stored SSL configuration for per-request override + ssl_config = self._ssl_verify + try: with map_aiohttp_exceptions(): response = await self._make_aiohttp_request( @@ -276,6 +295,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): timeout=timeout, proxy=proxy, sni_hostname=sni_hostname, + ssl_verify=ssl_config, ) except RuntimeError as e: # Handle the case where session was closed between our check and actual use @@ -296,6 +316,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): timeout=timeout, proxy=proxy, sni_hostname=sni_hostname, + ssl_verify=ssl_config, ) else: # Re-raise if it's a different RuntimeError diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index ac9dd5998e2..95f411c397c 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -846,6 +846,16 @@ class AsyncHTTPHandler: if str_to_bool(os.getenv("AIOHTTP_TRUST_ENV", "False")) is True: trust_env = True + ######################################################### + # Determine SSL config to pass to transport for per-request override + # This ensures ssl_verify works even with shared sessions + ######################################################### + ssl_for_transport: Optional[Union[bool, ssl.SSLContext]] = None + if ssl_context is not None: + ssl_for_transport = ssl_context + elif ssl_verify is False: + ssl_for_transport = False + verbose_logger.debug("Creating AiohttpTransport...") # Use shared session if provided and valid @@ -853,7 +863,10 @@ class AsyncHTTPHandler: verbose_logger.debug( f"SHARED SESSION: Reusing existing ClientSession (ID: {id(shared_session)})" ) - return LiteLLMAiohttpTransport(client=shared_session) + return LiteLLMAiohttpTransport( + client=shared_session, + ssl_verify=ssl_for_transport, + ) # Create new session only if none provided or existing one is invalid verbose_logger.debug( @@ -877,6 +890,7 @@ class AsyncHTTPHandler: connector=TCPConnector(**transport_connector_kwargs), trust_env=trust_env, ), + ssl_verify=ssl_for_transport, ) @staticmethod diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 95db8ec64b3..a97ebd8e74c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -81,6 +81,9 @@ from litellm.types.llms.anthropic_skills import ( ListSkillsResponse, Skill, ) +from litellm.anthropic_beta_headers_manager import ( + update_headers_with_filtered_beta, + ) from litellm.types.llms.openai import ( CreateBatchRequest, CreateFileRequest, @@ -1858,6 +1861,10 @@ class BaseLLMHTTPHandler: api_key=api_key, api_base=api_base, ) + + headers = update_headers_with_filtered_beta( + headers=headers, provider=custom_llm_provider + ) logging_obj.update_environment_variables( model=model, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index e9ae94307d4..7c2a9569c58 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -60,6 +60,38 @@ from ...anthropic.chat.transformation import AnthropicConfig from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import DatabricksBase, DatabricksException +def _sanitize_empty_content(message_dict: dict[str, Any]) -> None: + """ + Remove or filter content so empty text blocks are not sent. + Databricks Model Serving uses Anthropic Messages API spec and rejects empty text blocks. + """ + content = message_dict.get("content") + if content is None: + message_dict.pop("content", None) + return + if isinstance(content, str): + if not content.strip(): + message_dict.pop("content") + return + if isinstance(content, list): + if not content: + message_dict.pop("content") + return + filtered = [ + block + for block in content + if not ( + isinstance(block, dict) + and block.get("type") == "text" + and not (block.get("text") or "").strip() + ) + ] + if not filtered: + message_dict.pop("content") + else: + message_dict["content"] = filtered + + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -350,6 +382,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): # Move message-level cache_control into a content block when content is a string. if "cache_control" in _message and isinstance(_message.get("content"), str): _message = self._move_cache_control_into_string_content_block(_message) + _sanitize_empty_content(cast(dict[str, Any], _message)) new_messages.append(_message) if is_async: diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index e66394ae5f5..1c22602b483 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -838,6 +838,15 @@ class OCIChatConfig(BaseConfig): if not user_messages: raise Exception("No user message found for Cohere model") + # Extract system messages into preambleOverride + system_messages = [msg for msg in messages if msg.get("role") == "system"] + preamble_override = None + if system_messages: + preamble = "\n".join( + self._extract_text_content(msg["content"]) for msg in system_messages + ) + if preamble: + preamble_override = preamble # Create Cohere-specific chat request optional_cohere_params = self._get_optional_params(OCIVendors.COHERE, optional_params) @@ -845,6 +854,7 @@ class OCIChatConfig(BaseConfig): apiFormat="COHERE", message=self._extract_text_content(user_messages[-1]["content"]), chatHistory=self.adapt_messages_to_cohere_standard(messages), + preambleOverride=preamble_override, **optional_cohere_params ) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 8c98cc54050..bc5aa654aad 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -502,13 +502,12 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): reasoning_content: Optional[str] = None content: Optional[str] = None if chunk["message"].get("thinking") is not None: - if self.started_reasoning_content is False: - reasoning_content = chunk["message"].get("thinking") - self.started_reasoning_content = True - elif self.finished_reasoning_content is False: - reasoning_content = chunk["message"].get("thinking") - self.finished_reasoning_content = True + reasoning_content = chunk["message"].get("thinking") + self.started_reasoning_content = True elif chunk["message"].get("content") is not None: + if self.started_reasoning_content and not self.finished_reasoning_content: + self.finished_reasoning_content = True + message_content = chunk["message"].get("content") if "" in message_content: message_content = message_content.replace("", "") diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 6cc09dafc2f..16368907070 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -20,12 +20,12 @@ from typing import ( import httpx import litellm +from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _extract_reasoning_content, _handle_invalid_parallel_tool_calls, _should_convert_tool_call_to_json_mode, ) -from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.prompt_templates.common_utils import get_tool_call_names from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, @@ -161,6 +161,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "web_search_options", "service_tier", "safety_identifier", + "prompt_cache_key", ] # works across all models model_specific_params = [] diff --git a/litellm/llms/perplexity/responses/__init__.py b/litellm/llms/perplexity/responses/__init__.py new file mode 100644 index 00000000000..9bdf810e839 --- /dev/null +++ b/litellm/llms/perplexity/responses/__init__.py @@ -0,0 +1,7 @@ +""" +Perplexity Agentic Research API (Responses API) module +""" + +from .transformation import PerplexityResponsesConfig + +__all__ = ["PerplexityResponsesConfig"] diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py new file mode 100644 index 00000000000..178e76ea970 --- /dev/null +++ b/litellm/llms/perplexity/responses/transformation.py @@ -0,0 +1,409 @@ +""" +Transformation logic for Perplexity Agentic Research API (Responses API) + +This module handles the translation between OpenAI's Responses API format +and Perplexity's Responses API format, which supports: +- Third-party model access (OpenAI, Anthropic, Google, xAI, etc.) +- Presets for optimized configurations +- Web search and URL fetching tools +- Reasoning effort control +- Instructions parameter for system-level guidance +""" + +from typing import Any, Dict, List, Optional, Union + +import httpx + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +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 +from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseInputParam, + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + + +class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): + """ + Configuration for Perplexity Agentic Research API (Responses API) + + + Reference: https://docs.perplexity.ai/agentic-research/quickstart + """ + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.PERPLEXITY + + def get_supported_openai_params(self, model: str) -> list: + """ + Perplexity Responses API supports a different set of parameters + + Ref: https://docs.perplexity.ai/api-reference/responses-post + """ + return [ + "max_output_tokens", + "stream", + "temperature", + "top_p", + "tools", + "reasoning", + "preset", + "instructions", + "models", # Model fallback support + ] + + def validate_environment( + self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] + ) -> dict: + """Validate environment and set up headers""" + # Get API key from environment + api_key = ( + get_secret_str("PERPLEXITYAI_API_KEY") + or get_secret_str("PERPLEXITY_API_KEY") + ) + + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + headers["Content-Type"] = "application/json" + + return headers + + def get_complete_url( + self, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + """Get the complete URL for the Perplexity Responses API""" + if api_base is None: + api_base = get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" + + # Ensure api_base doesn't end with a slash + api_base = api_base.rstrip("/") + + # Add the responses endpoint + return f"{api_base}/v1/responses" + + def map_openai_params( + self, + response_api_optional_params: ResponsesAPIOptionalRequestParams, + model: str, + drop_params: bool, + ) -> Dict: + """ + Map OpenAI Responses API parameters to Perplexity format + + Key differences: + - Supports 'preset' parameter for predefined configurations + - Supports 'instructions' parameter for system-level guidance + - Tools are specified differently (web_search, fetch_url) + """ + mapped_params: Dict[str, Any] = {} + + # Map standard parameters + if response_api_optional_params.get("max_output_tokens"): + mapped_params["max_output_tokens"] = response_api_optional_params["max_output_tokens"] + + if response_api_optional_params.get("temperature"): + mapped_params["temperature"] = response_api_optional_params["temperature"] + + if response_api_optional_params.get("top_p"): + mapped_params["top_p"] = response_api_optional_params["top_p"] + + if response_api_optional_params.get("stream"): + mapped_params["stream"] = response_api_optional_params["stream"] + + if response_api_optional_params.get("stream_options"): + mapped_params["stream_options"] = response_api_optional_params["stream_options"] + + # Map Perplexity-specific parameters (using .get() with Any dict access) + preset = response_api_optional_params.get("preset") # type: ignore + if preset: + mapped_params["preset"] = preset + + instructions = response_api_optional_params.get("instructions") # type: ignore + if instructions: + mapped_params["instructions"] = instructions + + if response_api_optional_params.get("reasoning"): + mapped_params["reasoning"] = response_api_optional_params["reasoning"] + + tools = response_api_optional_params.get("tools") + if tools: + # Convert tools to list of dicts for transformation + tools_list = [dict(tool) if hasattr(tool, '__dict__') else tool for tool in tools] # type: ignore + mapped_params["tools"] = self._transform_tools(tools_list) # type: ignore + + return mapped_params + + def _transform_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Transform tools to Perplexity format + + Perplexity supports: + - web_search: Performs web searches + - fetch_url: Fetches content from URLs + """ + perplexity_tools = [] + + for tool in tools: + if isinstance(tool, dict): + tool_type = tool.get("type") + + # Direct Perplexity tool format + if tool_type in ["web_search", "fetch_url"]: + perplexity_tools.append(tool) + + # OpenAI function format - try to map to Perplexity tools + elif tool_type == "function": + function = tool.get("function", {}) + function_name = function.get("name", "") + + if function_name == "web_search" or "search" in function_name.lower(): + perplexity_tools.append({"type": "web_search"}) + elif function_name == "fetch_url" or "fetch" in function_name.lower(): + perplexity_tools.append({"type": "fetch_url"}) + + return perplexity_tools + + def transform_responses_api_request( + self, + model: str, + input: Union[str, ResponseInputParam], + response_api_optional_request_params: Dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Dict: + """ + Transform request to Perplexity Responses API format + """ + # Check if the model is a preset (format: preset/preset-name) + if model.startswith("preset/"): + preset_name = model.replace("preset/", "") + data = { + "preset": preset_name, + "input": self._format_input(input), + } + # Check if preset is explicitly provided in params + elif response_api_optional_request_params.get("preset"): + data = { + "preset": response_api_optional_request_params.pop("preset"), + "input": self._format_input(input), + } + else: + # Full request format for third-party models + data = { + "model": model, + "input": self._format_input(input), + } + + # Add all optional parameters + for key, value in response_api_optional_request_params.items(): + data[key] = value + + return data + + def _format_input(self, input: Union[str, ResponseInputParam]) -> Union[str, List[Dict[str, Any]]]: + """ + Format input for Perplexity Responses API + + The API accepts either: + - A simple string for single-turn queries + - An array of message objects for multi-turn conversations + """ + if isinstance(input, str): + return input + + # Handle ResponseInputParam format + if isinstance(input, list): + formatted_messages = [] + for item in input: + if isinstance(item, dict): + formatted_message = { + "type": "message", + "role": item.get("role"), + "content": item.get("content", ""), + } + formatted_messages.append(formatted_message) + return formatted_messages + + return str(input) + + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIResponse: + """ + Transform Perplexity Responses API response to OpenAI Responses API format + """ + try: + raw_response_json = raw_response.json() + except Exception as e: + raise BaseLLMException( + status_code=raw_response.status_code, + message=f"Failed to parse response: {str(e)}", + ) + + # Check for error status + status = raw_response_json.get("status") + if status == "failed": + error = raw_response_json.get("error", {}) + error_message = error.get("message", "Unknown error") + raise BaseLLMException( + status_code=raw_response.status_code, + message=error_message, + ) + + # Transform usage to handle Perplexity's cost structure + usage_data = raw_response_json.get("usage", {}) + transformed_usage_dict = self._transform_usage(usage_data) + + # Convert usage dict to ResponseAPIUsage object + usage_obj = ResponseAPIUsage(**transformed_usage_dict) if transformed_usage_dict else None + + # Map Perplexity response to OpenAI Responses API format + response = ResponsesAPIResponse( + id=raw_response_json.get("id", ""), + object="response", + created_at=raw_response_json.get("created_at", 0), + status=raw_response_json.get("status", "completed"), + model=raw_response_json.get("model", model), + output=raw_response_json.get("output", []), + usage=usage_obj, + ) + + return response + + def _transform_usage(self, usage_data: Dict[str, Any]) -> Dict[str, Any]: + """ + Transform Perplexity usage data to OpenAI format + + Perplexity returns: + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost": { + "currency": "USD", + "input_cost": 0.0001, + "output_cost": 0.0002, + "total_cost": 0.0003 + } + } + + OpenAI expects: + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost": 0.0003 + } + """ + transformed = { + "input_tokens": usage_data.get("input_tokens", 0), + "output_tokens": usage_data.get("output_tokens", 0), + "total_tokens": usage_data.get("total_tokens", 0), + } + + # Transform cost from Perplexity format (dict) to OpenAI format (float) + cost_obj = usage_data.get("cost") + if isinstance(cost_obj, dict) and "total_cost" in cost_obj: + transformed["cost"] = cost_obj["total_cost"] + verbose_logger.debug( + "Transformed Perplexity cost object to float: %s -> %s", + cost_obj, + cost_obj["total_cost"] + ) + elif cost_obj is not None: + # If cost is already a float/number, use it as-is + transformed["cost"] = cost_obj + + # Add input_tokens_details if present + if "input_tokens_details" in usage_data: + transformed["input_tokens_details"] = usage_data["input_tokens_details"] + + # Add output_tokens_details if present + if "output_tokens_details" in usage_data: + transformed["output_tokens_details"] = usage_data["output_tokens_details"] + + return transformed + + def transform_streaming_response( + self, + model: str, + parsed_chunk: dict, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIStreamingResponse: + """ + Transform a parsed streaming response chunk into a ResponsesAPIStreamingResponse + """ + # Get the event type from the chunk + verbose_logger.debug("Raw Perplexity Chunk=%s", parsed_chunk) + event_type = str(parsed_chunk.get("type")) + event_pydantic_model = PerplexityResponsesConfig.get_event_model_class( + event_type=event_type + ) + + # Transform Perplexity-specific fields to OpenAI format + parsed_chunk = self._transform_perplexity_chunk(parsed_chunk) + + # Defensive: Handle error.code being null (similar to OpenAI implementation) + try: + error_obj = parsed_chunk.get("error") + if isinstance(error_obj, dict) and error_obj.get("code") is None: + # Preserve other fields, but ensure `code` is a non-null string + parsed_chunk = dict(parsed_chunk) + parsed_chunk["error"] = dict(error_obj) + parsed_chunk["error"]["code"] = "unknown_error" + except Exception: + # If anything unexpected happens here, fall back to attempting + # instantiation and let higher-level handlers manage errors. + verbose_logger.debug("Failed to coalesce error.code in parsed_chunk") + + return event_pydantic_model(**parsed_chunk) + + def _transform_perplexity_chunk(self, chunk: dict) -> dict: + """ + Transform Perplexity-specific fields in a streaming chunk to OpenAI format. + + This handles: + - Converting Perplexity's cost object to a simple float + """ + # Make a copy to avoid modifying the original + chunk = dict(chunk) + + # Transform usage.cost from Perplexity format to OpenAI format + # Perplexity: {"currency": "USD", "input_cost": 0.0001, "output_cost": 0.0002, "total_cost": 0.0003} + # OpenAI: 0.0003 (just the total_cost as a float) + try: + response_obj = chunk.get("response") + if isinstance(response_obj, dict): + usage_obj = response_obj.get("usage") + if isinstance(usage_obj, dict): + cost_obj = usage_obj.get("cost") + if isinstance(cost_obj, dict) and "total_cost" in cost_obj: + # Replace the cost object with just the total_cost value + chunk = dict(chunk) + chunk["response"] = dict(response_obj) + chunk["response"]["usage"] = dict(usage_obj) + chunk["response"]["usage"]["cost"] = cost_obj["total_cost"] + verbose_logger.debug( + "Transformed Perplexity cost object to float: %s -> %s", + cost_obj, + cost_obj["total_cost"] + ) + except Exception as e: + # If transformation fails, log and continue with original chunk + verbose_logger.debug("Failed to transform Perplexity cost object: %s", e) + + return chunk diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index a0e2ddf5e98..02b69b94d94 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1,4 +1,5 @@ import re +from copy import deepcopy from enum import Enum from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, get_type_hints @@ -684,7 +685,7 @@ def convert_anyof_null_to_nullable(schema, depth=0): if anyof is not None: contains_null = False for atype in anyof: - if atype == {"type": "null"}: + if isinstance(atype, dict) and atype.get("type") == "null": # remove null type anyof.remove(atype) contains_null = True @@ -801,8 +802,38 @@ def _convert_schema_types(schema, depth=0): if "type" in schema: type_val = schema["type"] if isinstance(type_val, list) and len(type_val) > 1: - # Convert ["string", "number"] -> {"anyOf": [{"type": "STRING"}, {"type": "NUMBER"}]} - schema["anyOf"] = [{"type": t} for t in type_val if isinstance(t, str)] + # Convert type arrays to anyOf format + # Fields that are specific to object/array types and should move into anyOf + type_specific_fields = {"properties", "required", "additionalProperties", "items", "minItems", "maxItems", "minProperties", "maxProperties"} + + any_of: List[Dict[str, Any]] = [] + for t in type_val: + if not isinstance(t, str): + continue + if t == "null": + # Keep null entry minimal so we can strip it later. + any_of.append({"type": "null"}) + continue + + # For object/array types, include type-specific fields + if t in ("object", "array"): + item_schema = {"type": t} + # Move type-specific fields into this anyOf item + for field in type_specific_fields: + if field in schema: + item_schema[field] = deepcopy(schema[field]) + any_of.append(item_schema) + else: + # For primitive types, only include the type + any_of.append({"type": t}) + + # Remove type-specific fields from parent if we moved them into anyOf + has_object_or_array = any(t in ("object", "array") for t in type_val if isinstance(t, str)) + if has_object_or_array: + for field in type_specific_fields: + schema.pop(field, None) + + schema["anyOf"] = any_of schema.pop("type") elif isinstance(type_val, list) and len(type_val) == 1: schema["type"] = type_val[0] diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 3004f39b973..00b461dcda0 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -437,6 +437,27 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 else: assistant_content.append(PartType(text=assistant_text)) # type: ignore + ## HANDLE ASSISTANT IMAGES FIELD + # Process images field if present (for generated images from assistant) + assistant_images = assistant_msg.get("images") + if assistant_images is not None and isinstance(assistant_images, list): + for image_item in assistant_images: + if isinstance(image_item, dict): + image_url_obj = image_item.get("image_url") + if isinstance(image_url_obj, dict): + assistant_image_url = image_url_obj.get("url") + format = image_url_obj.get("format") + detail = image_url_obj.get("detail") + media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) + if assistant_image_url: + _part = _process_gemini_media( + image_url=assistant_image_url, + format=format, + media_resolution_enum=media_resolution_enum, + model=model, + ) + assistant_content.append(_part) + ## HANDLE ASSISTANT FUNCTION CALL if ( assistant_msg.get("tool_calls", []) is not None 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 5a09168282d..54c3f9e0474 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,8 +1,5 @@ from typing import Any, Dict, List, Optional, Tuple -from litellm.anthropic_beta_headers_manager import ( - update_headers_with_filtered_beta, -) from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, @@ -105,12 +102,6 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert if beta_values: headers["anthropic-beta"] = ",".join(beta_values) - # Filter out unsupported beta headers for Vertex AI - headers = update_headers_with_filtered_beta( - headers=headers, - provider="vertex_ai", - ) - return headers, api_base def get_complete_url( 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 3842159fd7b..c6914ac3d6b 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 @@ -107,6 +107,11 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): vertex_project = self.get_vertex_ai_project(litellm_params) vertex_location = self.get_vertex_ai_location(litellm_params) + # Map empty location/cluade models to a supported region for count-tokens endpoint + # https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens + if not vertex_location or "claude" in model.lower(): + vertex_location = "us-central1" + # Get access token and resolved project ID access_token, project_id = await self._ensure_access_token_async( credentials=vertex_credentials, @@ -118,7 +123,7 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): endpoint_url = self._build_count_tokens_endpoint( model=model, project_id=project_id, - vertex_location=vertex_location or "us-central1", + vertex_location=vertex_location, api_base=litellm_params.get("api_base"), ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 815d29c7964..f6edcf7efd0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5848,6 +5848,19 @@ "output_cost_per_token": 7e-07, "supports_tool_choice": true }, + "azure_ai/kimi-k2.5": { + "input_cost_per_token": 6e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/kimi-k2-5-now-in-microsoft-foundry/4492321", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure_ai/ministral-3b": { "input_cost_per_token": 4e-08, "litellm_provider": "azure_ai", @@ -6091,6 +6104,39 @@ "output_cost_per_token": 2.4e-05, "supports_tool_choice": true }, + "bedrock/ap-northeast-1/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 7.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.03e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "bedrock/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 7.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.03e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "bedrock/moonshotai.kimi-k2.5": { + "input_cost_per_token": 7.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.03e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/ap-south-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 3.18e-06, "litellm_provider": "bedrock", @@ -6109,6 +6155,17 @@ "mode": "chat", "output_cost_per_token": 7.2e-07 }, + "bedrock/ap-south-1/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 7.1e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.94e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/ca-central-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 3.05e-06, "litellm_provider": "bedrock", @@ -6314,6 +6371,17 @@ "mode": "chat", "output_cost_per_token": 1.01e-06 }, + "bedrock/sa-east-1/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 7.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.03e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/us-east-1/1-month-commitment/anthropic.claude-instant-v1": { "input_cost_per_second": 0.011, "litellm_provider": "bedrock", @@ -6450,6 +6518,28 @@ "output_cost_per_token": 7e-07, "supports_tool_choice": true }, + "bedrock/us-east-1/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 6e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "bedrock/us-east-2/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 6e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/us-gov-east-1/amazon.nova-pro-v1:0": { "input_cost_per_token": 9.6e-07, "litellm_provider": "bedrock", @@ -6856,6 +6946,17 @@ "output_cost_per_token": 7e-07, "supports_tool_choice": true }, + "bedrock/us-west-2/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 6e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, "cache_read_input_token_cost": 8e-08, @@ -8945,6 +9046,43 @@ } ] }, + "dashscope/qwen3-max": { + "litellm_provider": "dashscope", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, "dashscope/qwq-plus": { "input_cost_per_token": 8e-07, "litellm_provider": "dashscope", @@ -10618,14 +10756,22 @@ "input_cost_per_token": 2.8e-07, "input_cost_per_token_cache_hit": 2.8e-08, "litellm_provider": "deepseek", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 4.2e-07, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, "supports_tool_choice": true }, "deepseek/deepseek-coder": { @@ -10662,16 +10808,24 @@ "input_cost_per_token": 2.8e-07, "input_cost_per_token_cache_hit": 2.8e-08, "litellm_provider": "deepseek", - "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 4.2e-07, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_assistant_prefill": true, - "supports_function_calling": true, + "supports_function_calling": false, + "supports_native_streaming": true, + "supports_parallel_function_calling": false, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false }, "deepseek/deepseek-v3": { "cache_creation_input_token_cost": 0.0, @@ -25614,6 +25768,66 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "perplexity/preset/pro-search": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_preset": true + }, + "perplexity/openai/gpt-4o": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/openai/gpt-4o-mini": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/openai/gpt-5.2": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": true + }, + "perplexity/anthropic/claude-3-5-sonnet-20241022": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/anthropic/claude-3-5-haiku-20241022": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/google/gemini-2.0-flash-exp": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/google/gemini-2.0-flash-thinking-exp": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": true + }, + "perplexity/xai/grok-2-1212": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/xai/grok-2-vision-1212": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, "publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": { "input_cost_per_token": 0.0, "litellm_provider": "publicai", diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 8b052dd0da1..b0ae03c94d3 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,6 +1,6 @@ import json from typing import Optional -from urllib.parse import urlencode, urlparse, urlunparse +from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse @@ -16,6 +16,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.utils import get_server_root_path +from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer router = APIRouter( @@ -125,6 +126,29 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +def _resolve_oauth2_server_for_root_endpoints( + client_ip: Optional[str] = None, +) -> Optional[MCPServer]: + """ + Resolve the MCP server for root-level OAuth endpoints (no server name in path). + + When the MCP SDK hits root-level endpoints like /register, /authorize, /token + without a server name prefix, we try to find the right server automatically. + Returns the server if exactly one OAuth2 server is configured, else None. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + registry = global_mcp_server_manager.get_filtered_registry(client_ip=client_ip) + oauth2_servers = [ + s for s in registry.values() if s.auth_type == MCPAuth.oauth2 + ] + if len(oauth2_servers) == 1: + return oauth2_servers[0] + return None + + async def authorize_with_server( request: Request, mcp_server: MCPServer, @@ -170,7 +194,13 @@ async def authorize_with_server( if code_challenge_method: params["code_challenge_method"] = code_challenge_method - return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}") + parsed_auth_url = urlparse(mcp_server.authorization_url) + existing_params = dict(parse_qsl(parsed_auth_url.query)) + existing_params.update(params) + final_url = urlunparse( + parsed_auth_url._replace(query=urlencode(existing_params)) + ) + return RedirectResponse(final_url) async def exchange_token_with_server( @@ -305,6 +335,8 @@ async def authorize( mcp_server = global_mcp_server_manager.get_mcp_server_by_name( lookup_name, client_ip=client_ip ) + if mcp_server is None and mcp_server_name is None: + mcp_server = _resolve_oauth2_server_for_root_endpoints() if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") return await authorize_with_server( @@ -350,6 +382,8 @@ async def token_endpoint( mcp_server = global_mcp_server_manager.get_mcp_server_by_name( lookup_name, client_ip=client_ip ) + if mcp_server is None and mcp_server_name is None: + mcp_server = _resolve_oauth2_server_for_root_endpoints() if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") return await exchange_token_with_server( @@ -430,6 +464,13 @@ def _build_oauth_protected_resource_response( ) request_base_url = get_request_base_url(request) + + # When no server name provided, try to resolve the single OAuth2 server + if mcp_server_name is None: + resolved = _resolve_oauth2_server_for_root_endpoints() + if resolved: + mcp_server_name = resolved.server_name or resolved.name + mcp_server: Optional[MCPServer] = None if mcp_server_name: client_ip = IPAddressUtils.get_mcp_client_ip(request) @@ -535,6 +576,12 @@ def _build_oauth_authorization_server_response( request_base_url = get_request_base_url(request) + # When no server name provided, try to resolve the single OAuth2 server + if mcp_server_name is None: + resolved = _resolve_oauth2_server_for_root_endpoints() + if resolved: + mcp_server_name = resolved.server_name or resolved.name + authorization_endpoint = ( f"{request_base_url}/{mcp_server_name}/authorize" if mcp_server_name @@ -640,6 +687,19 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non "redirect_uris": [f"{request_base_url}/callback"], } if not mcp_server_name: + resolved = _resolve_oauth2_server_for_root_endpoints() + if resolved: + return await register_client_with_server( + request=request, + mcp_server=resolved, + client_name=data.get("client_name", ""), + grant_types=data.get("grant_types", []), + response_types=data.get("response_types", []), + token_endpoint_auth_method=data.get( + "token_endpoint_auth_method", "" + ), + fallback_client_id=resolved.server_name or resolved.name, + ) return dummy_return client_ip = IPAddressUtils.get_mcp_client_ip(request) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8878c52b077..fd251488db4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -339,7 +339,7 @@ class MCPServerManager: verbose_logger.info( f"Loading OpenAPI spec from {spec_path} for server {server_name}" ) - self._register_openapi_tools( + await self._register_openapi_tools( spec_path=spec_path, server=new_server, base_url=server_config.get("url", ""), @@ -351,7 +351,9 @@ class MCPServerManager: self.initialize_tool_name_to_mcp_server_name_mapping() - def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_url: str): + async def _register_openapi_tools( + self, spec_path: str, server: MCPServer, base_url: str + ): """ Register tools from an OpenAPI specification for a given server. @@ -373,15 +375,15 @@ class MCPServerManager: get_base_url as get_openapi_base_url, ) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( - load_openapi_spec, + load_openapi_spec_async, ) from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) try: - # Load OpenAPI spec - spec = load_openapi_spec(spec_path) + # Load OpenAPI spec (async to avoid "called from within a running event loop") + spec = await load_openapi_spec_async(spec_path) # Use base_url from config if provided, otherwise extract from spec if not base_url: diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index b635f15ed09..deb0b4f9549 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -3,6 +3,8 @@ This module is used to generate MCP tools from OpenAPI specs. """ import json +import asyncio +import os from pathlib import PurePosixPath from typing import Any, Dict, Optional from urllib.parse import quote @@ -45,8 +47,36 @@ def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: def load_openapi_spec(filepath: str) -> Dict[str, Any]: - """Load OpenAPI specification from JSON file.""" - with open(filepath, "r") as f: + """ + Sync wrapper. For URL specs, use the shared/custom MCP httpx client. + """ + try: + # If we're already inside an event loop, prefer the async function. + asyncio.get_running_loop() + raise RuntimeError( + "load_openapi_spec() was called from within a running event loop. " + "Use 'await load_openapi_spec_async(...)' instead." + ) + except RuntimeError as e: + # "no running event loop" is fine; other RuntimeErrors we re-raise + if "no running event loop" not in str(e).lower(): + raise + return asyncio.run(load_openapi_spec_async(filepath)) + +async def load_openapi_spec_async(filepath: str) -> Dict[str, Any]: + if filepath.startswith("http://") or filepath.startswith("https://"): + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) + # NOTE: do not close shared client if get_async_httpx_client returns a shared singleton. + # If it returns a new client each time, consider wrapping it in an async context manager. + r = await client.get(filepath) + r.raise_for_status() + return r.json() + + # fallback: local file + # Local filesystem path + if not os.path.exists(filepath): + raise FileNotFoundError(f"OpenAPI spec not found at {filepath}") + with open(filepath, "r", encoding="utf-8") as f: return json.load(f) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 7b0253992e0..aed81afd254 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,6 +1,6 @@ import importlib from datetime import datetime -from typing import Dict, List, Optional, Union +from typing import Any, Awaitable, Callable, Dict, List, Optional, Union from fastapi import APIRouter, Depends, HTTPException, Query, Request @@ -528,24 +528,50 @@ if MCP_AVAILABLE: NewMCPServerRequest, ) + def _extract_credentials( + request: NewMCPServerRequest, + ) -> tuple: + """ + Extract OAuth credentials from the nested ``request.credentials`` dict. + + Returns: + (client_id, client_secret, scopes) — any value may be ``None``. + """ + creds = request.credentials if isinstance(request.credentials, dict) else {} + client_id: Optional[str] = creds.get("client_id") + client_secret: Optional[str] = creds.get("client_secret") + scopes_raw = creds.get("scopes") + scopes: Optional[List[str]] = scopes_raw if isinstance(scopes_raw, list) else None + return client_id, client_secret, scopes + async def _execute_with_mcp_client( request: NewMCPServerRequest, - operation, + operation: Callable[..., Awaitable[Any]], mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, - ): + ) -> dict: """ - Common helper to create MCP client, execute operation, and ensure proper cleanup. + Create a temporary MCP client from *request*, run *operation*, and return the result. + + For M2M OAuth servers (those with ``client_id``, ``client_secret``, and + ``token_url``), the incoming ``oauth2_headers`` are dropped so that + ``resolve_mcp_auth`` can auto-fetch a token via ``client_credentials``. Args: - request: MCP server configuration - operation: Async function that takes a client and returns the operation result + request: MCP server configuration submitted by the UI. + operation: Async callable that receives the created client and returns a result dict. + mcp_auth_header: Pre-resolved credential header (API-key / bearer token). + oauth2_headers: Headers extracted from the incoming request (may contain the + litellm API key — must NOT be forwarded for M2M servers). + raw_headers: Raw request headers forwarded for stdio env construction. Returns: - Operation result or error response + The dict returned by *operation*, or an error dict on failure. """ try: + client_id, client_secret, scopes = _extract_credentials(request) + server_model = MCPServer( server_id=request.server_id or "", name=request.alias or request.server_name or "", @@ -557,14 +583,26 @@ if MCP_AVAILABLE: args=request.args, env=request.env, static_headers=request.static_headers, + client_id=client_id, + client_secret=client_secret, + token_url=request.token_url, + scopes=scopes, + authorization_url=request.authorization_url, + registration_url=request.registration_url, ) stdio_env = global_mcp_server_manager._build_stdio_env( server_model, raw_headers ) + # For M2M OAuth servers, drop the incoming Authorization header so that + # resolve_mcp_auth can auto-fetch a token via client_credentials. + effective_oauth2_headers = ( + None if server_model.has_client_credentials else oauth2_headers + ) + merged_headers = merge_mcp_headers( - extra_headers=oauth2_headers, + extra_headers=effective_oauth2_headers, static_headers=request.static_headers, ) @@ -577,11 +615,14 @@ if MCP_AVAILABLE: return await operation(client) - except Exception as e: - verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True) + except (KeyboardInterrupt, SystemExit): + raise + except BaseException as e: + verbose_logger.error("Error in MCP operation: %s", e, exc_info=True) return { "status": "error", - "message": "An internal error has occurred while testing the MCP server.", + "error": True, + "message": "Failed to connect to MCP server. Check proxy logs for details.", } @router.post("/test/connection", dependencies=[Depends(user_api_key_auth)]) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 890c4ae8fb2..58cd8c99e7b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -31,6 +31,9 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) +from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, +) from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, @@ -1972,7 +1975,7 @@ if MCP_AVAILABLE: ) if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: request = StarletteRequest(scope) - base_url = str(request.base_url).rstrip("/") + base_url = get_request_base_url(request) authorization_uri = ( f"Bearer authorization_uri=" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bb3fd748ef3..87ff4a66e08 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -633,6 +633,9 @@ class LiteLLMRoutes(enum.Enum): "/model/{model_id}/update", "/prompt/list", "/prompt/info", + # Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges + "/invitation/new", + "/invitation/delete", ] # routes that manage their own allowed/disallowed logic ## Org Admin Routes ## @@ -1911,6 +1914,10 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase): default=None, description="Guardrails configuration for this passthrough endpoint. Dict keys are guardrail names, values are optional settings for field targeting. When set, all org/team/key level guardrails will also execute. Defaults to None (no guardrails execute).", ) + is_from_config: bool = Field( + default=False, + description="True if this endpoint is defined in the config file, False if from DB. Config-defined endpoints cannot be edited via the UI.", + ) class PassThroughEndpointResponse(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 05eeab3f611..42f10ff8598 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -8,6 +8,7 @@ Returns a UserAPIKeyAuth object if the API key is valid """ import asyncio +import re import secrets from datetime import datetime, timezone from typing import List, Optional, Tuple, cast @@ -115,6 +116,18 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str: api_key = api_key.replace("Basic ", "") # handle langfuse input elif api_key.startswith("bearer "): api_key = api_key.replace("bearer ", "") + elif api_key.startswith("AWS4-HMAC-SHA256"): + # Handle AWS Signature V4 format from LangChain + # Format: AWS4-HMAC-SHA256 Credential=Bearer sk-12345/date/region/service/aws4_request, SignedHeaders=..., Signature=... + # Extract the Bearer token from the Credential field + match = re.search(r'Credential=Bearer\s+([^/\s,]+)', api_key) + if match: + api_key = match.group(1) + else: + # If no Bearer token found in Credential, try to extract just the credential value + match = re.search(r'Credential=([^/\s,]+)', api_key) + if match: + api_key = match.group(1) return api_key @@ -128,6 +141,20 @@ def _get_bearer_token( api_key = api_key.replace("Basic ", "") # handle langfuse input elif api_key.startswith("bearer "): api_key = api_key.replace("bearer ", "") + elif api_key.startswith("AWS4-HMAC-SHA256"): + # Handle AWS Signature V4 format from LangChain + # Format: AWS4-HMAC-SHA256 Credential=Bearer sk-12345/date/region/service/aws4_request, SignedHeaders=..., Signature=... + # Extract the Bearer token from the Credential field + match = re.search(r'Credential=Bearer\s+([^/\s,]+)', api_key) + if match: + api_key = match.group(1) + else: + # If no Bearer token found in Credential, try to extract just the credential value + match = re.search(r'Credential=([^/\s,]+)', api_key) + if match: + api_key = match.group(1) + else: + api_key = "" else: api_key = "" return api_key diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index faeca9b2aed..62ca6dc2ae2 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -394,6 +394,14 @@ def get_logging_caching_headers(request_data: Dict) -> Optional[Dict]: _metadata["applied_policies"] ) + if "policy_sources" in _metadata: + sources = _metadata["policy_sources"] + if isinstance(sources, dict) and sources: + # Use ';' as delimiter — matched_via reasons may contain commas + headers["x-litellm-policy-sources"] = "; ".join( + f"{name}={reason}" for name, reason in sources.items() + ) + if "semantic-similarity" in _metadata: headers["x-litellm-semantic-similarity"] = str(_metadata["semantic-similarity"]) @@ -441,6 +449,27 @@ def add_policy_to_applied_policies_header( request_data["metadata"] = _metadata +def add_policy_sources_to_metadata( + request_data: Dict, policy_sources: Dict[str, str] +): + """ + Store policy match reasons in metadata for x-litellm-policy-sources header. + + Args: + request_data: The request data dict + policy_sources: Map of policy_name -> matched_via reason + """ + if not policy_sources: + return + _metadata = request_data.get("metadata", None) or {} + existing = _metadata.get("policy_sources", {}) + if not isinstance(existing, dict): + existing = {} + existing.update(policy_sources) + _metadata["policy_sources"] = existing + request_data["metadata"] = _metadata + + def add_guardrail_response_to_standard_logging_object( litellm_logging_obj: Optional["LiteLLMLogging"], guardrail_response: StandardLoggingGuardrailInformation, diff --git a/litellm/proxy/custom_sso.py b/litellm/proxy/custom_sso.py index 210e9eea3dc..b2b028dfbe3 100644 --- a/litellm/proxy/custom_sso.py +++ b/litellm/proxy/custom_sso.py @@ -24,9 +24,12 @@ async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: print(f"userIDPInfo: {userIDPInfo}") # noqa if userIDPInfo.id is None: - raise ValueError( - f"No ID found for user. userIDPInfo.id is None {userIDPInfo}" - ) + raise ValueError(f"No ID found for user. userIDPInfo.id is None {userIDPInfo}") + + # Access extra fields from the IDP response (requires GENERIC_USER_EXTRA_ATTRIBUTES env var) + # Example: Set GENERIC_USER_EXTRA_ATTRIBUTES="group,NTID,domain" to capture these fields + # extra_fields = getattr(userIDPInfo, 'extra_fields', None) or {} + # user_groups = extra_fields.get("group", []) # check if user exists in litellm proxy DB _user_info = await user_info(user_id=userIDPInfo.id) diff --git a/litellm/proxy/db/db_transaction_queue/base_update_queue.py b/litellm/proxy/db/db_transaction_queue/base_update_queue.py index 202829b78b6..a5ec1c3eaf4 100644 --- a/litellm/proxy/db/db_transaction_queue/base_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/base_update_queue.py @@ -10,14 +10,18 @@ from litellm._service_logger import ServiceLogging service_logger_obj = ( ServiceLogging() ) # used for tracking metrics for In memory buffer, redis buffer, pod lock manager -from litellm.constants import MAX_IN_MEMORY_QUEUE_FLUSH_COUNT, MAX_SIZE_IN_MEMORY_QUEUE +from litellm.constants import ( + LITELLM_ASYNCIO_QUEUE_MAXSIZE, + MAX_IN_MEMORY_QUEUE_FLUSH_COUNT, + MAX_SIZE_IN_MEMORY_QUEUE, +) class BaseUpdateQueue: """Base class for in memory buffer for database transactions""" def __init__(self): - self.update_queue = asyncio.Queue() + self.update_queue = asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) self.MAX_SIZE_IN_MEMORY_QUEUE = MAX_SIZE_IN_MEMORY_QUEUE async def add_update(self, update): diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index c3074e641b2..5ba8fb13596 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -3,6 +3,7 @@ from copy import deepcopy from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.proxy._types import BaseDailySpendTransaction from litellm.proxy.db.db_transaction_queue.base_update_queue import ( BaseUpdateQueue, @@ -54,7 +55,7 @@ class DailySpendUpdateQueue(BaseUpdateQueue): def __init__(self): super().__init__() self.update_queue: asyncio.Queue[Dict[str, BaseDailySpendTransaction]] = ( - asyncio.Queue() + asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) ) async def add_update(self, update: Dict[str, BaseDailySpendTransaction]): diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 9b0449bb9ab..c96564252d0 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -2,6 +2,7 @@ import asyncio from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.proxy._types import ( DBSpendUpdateTransactions, Litellm_EntityType, @@ -21,7 +22,9 @@ class SpendUpdateQueue(BaseUpdateQueue): def __init__(self): super().__init__() - self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue() + self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue( + maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE + ) async def flush_and_get_aggregated_db_spend_update_transactions( self, diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 3cbca27ce0c..955067b486a 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -19,13 +19,15 @@ async def get_ui_config(): from litellm.proxy.utils import get_proxy_base_url, get_server_root_path auto_redirect_ui_login_to_sso = ( - os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "true").lower() == "true" + os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "false").lower() == "true" ) admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" + sso_configured = _has_user_setup_sso() return UiDiscoveryEndpoints( server_root_path=get_server_root_path(), proxy_base_url=get_proxy_base_url(), - auto_redirect_to_sso=_has_user_setup_sso() and auto_redirect_ui_login_to_sso, + auto_redirect_to_sso=sso_configured and auto_redirect_ui_login_to_sso, admin_ui_disabled=admin_ui_disabled, + sso_configured=sso_configured, ) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index a825ce22b25..07702252a10 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -15,6 +15,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.types.guardrails import ( + BaseLitellmParams, PII_ENTITY_CATEGORIES_MAP, ApplyGuardrailRequest, ApplyGuardrailResponse, @@ -150,6 +151,7 @@ async def list_guardrails_v2(): } ``` """ + from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client @@ -164,11 +166,29 @@ async def list_guardrails_v2(): guardrail_configs: List[GuardrailInfoResponse] = [] seen_guardrail_ids = set() for guardrail in guardrails: + litellm_params: Optional[Union[LitellmParams, dict]] = guardrail.get( + "litellm_params" + ) + litellm_params_dict = ( + litellm_params.model_dump(exclude_none=True) + if isinstance(litellm_params, LitellmParams) + else litellm_params + ) or {} + masked_litellm_params_dict = _get_masked_values( + litellm_params_dict, + unmasked_length=4, + number_of_asterisks=4, + ) + masked_litellm_params = ( + BaseLitellmParams(**masked_litellm_params_dict) + if masked_litellm_params_dict + else None + ) guardrail_configs.append( GuardrailInfoResponse( guardrail_id=guardrail.get("guardrail_id"), guardrail_name=guardrail.get("guardrail_name"), - litellm_params=guardrail.get("litellm_params"), + litellm_params=masked_litellm_params, guardrail_info=guardrail.get("guardrail_info"), created_at=guardrail.get("created_at"), updated_at=guardrail.get("updated_at"), @@ -182,11 +202,27 @@ async def list_guardrails_v2(): for guardrail in in_memory_guardrails: # only add guardrails that are not in DB guardrail list already if guardrail.get("guardrail_id") not in seen_guardrail_ids: + in_memory_litellm_params_raw = guardrail.get("litellm_params") + in_memory_litellm_params_dict = ( + in_memory_litellm_params_raw.model_dump(exclude_none=True) + if isinstance(in_memory_litellm_params_raw, LitellmParams) + else in_memory_litellm_params_raw + ) or {} + masked_in_memory_litellm_params = _get_masked_values( + in_memory_litellm_params_dict, + unmasked_length=4, + number_of_asterisks=4, + ) + masked_in_memory_litellm_params_typed = ( + BaseLitellmParams(**masked_in_memory_litellm_params) + if masked_in_memory_litellm_params + else None + ) guardrail_configs.append( GuardrailInfoResponse( guardrail_id=guardrail.get("guardrail_id"), guardrail_name=guardrail.get("guardrail_name"), - litellm_params=dict(guardrail.get("litellm_params") or {}), + litellm_params=masked_in_memory_litellm_params_typed, guardrail_info=dict(guardrail.get("guardrail_info") or {}), guardrail_definition_location="config", ) @@ -666,11 +702,16 @@ async def get_guardrail_info(guardrail_id: str): unmasked_length=4, number_of_asterisks=4, ) + masked_litellm_params = ( + BaseLitellmParams(**masked_litellm_params_dict) + if masked_litellm_params_dict + else None + ) return GuardrailInfoResponse( guardrail_id=result.get("guardrail_id"), guardrail_name=result.get("guardrail_name"), - litellm_params=masked_litellm_params_dict, + litellm_params=masked_litellm_params, guardrail_info=dict(result.get("guardrail_info") or {}), created_at=result.get("created_at"), updated_at=result.get("updated_at"), diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 528857f5dd1..6800dff55ac 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -461,12 +461,37 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): data=prepared_request.body, # type: ignore headers=prepared_request.headers, # type: ignore ) + except HTTPException: + # Propagate HTTPException (e.g. from non-200 path) as-is + raise except Exception as e: + # If this is an HTTP error with a response body (e.g. httpx.HTTPStatusError), + # extract the AWS error message and propagate it + response = getattr(e, "response", None) + if isinstance(response, httpx.Response): + try: + status_code, detail_message = ( + self._parse_bedrock_guardrail_error_response(response) + ) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={"error": detail_message}, + request_data=request_data or {}, + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=datetime.now().timestamp(), + duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, + ) + raise HTTPException( + status_code=status_code, detail=detail_message + ) from e + except HTTPException: + raise # Endpoint down, timeout, or other HTTP/network errors verbose_proxy_logger.error( "Bedrock AI: failed to make guardrail request: %s", str(e) ) - # Add guardrail information with failure status self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": str(e)}, @@ -477,7 +502,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): duration=(datetime.now() - start_time).total_seconds(), event_type=event_type, ) - # Re-raise the exception to maintain existing behavior raise ######################################################### @@ -509,11 +533,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): bedrock_guardrail_response ) else: + status_code, detail_message = self._parse_bedrock_guardrail_error_response( + httpx_response + ) verbose_proxy_logger.error( "Bedrock AI: error in response. Status code: %s, response: %s", httpx_response.status_code, httpx_response.text, ) + raise HTTPException(status_code=status_code, detail=detail_message) return bedrock_guardrail_response @@ -579,6 +607,34 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return "success" return "guardrail_failed_to_respond" + def _parse_bedrock_guardrail_error_response( + self, response: httpx.Response + ) -> Tuple[int, str]: + """ + Parse AWS Bedrock guardrail error response body to extract status code and message. + + AWS may return shapes like {"message": "..."} or {"error": {"message": "..."}}. + Returns (status_code, message) for use in HTTPException. + """ + status_code = response.status_code + message = "Bedrock guardrail request failed" + try: + body = response.json() + except Exception: + text = getattr(response, "text", None) or "" + if isinstance(text, str) and text.strip(): + return (status_code, text.strip()) + return (status_code, message) + if isinstance(body, dict): + if isinstance(body.get("message"), str): + return (status_code, body["message"]) + err = body.get("error") + if isinstance(err, dict) and isinstance(err.get("message"), str): + return (status_code, err["message"]) + if isinstance(err, str): + return (status_code, err) + return (status_code, message) + def _get_http_exception_for_blocked_guardrail( self, response: BedrockGuardrailResponse ) -> Union[HTTPException, GuardrailInterventionNormalStringError]: @@ -739,9 +795,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[ - Union[BedrockGuardrailResponse, str] - ] = None + bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( + None + ) try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data @@ -811,9 +867,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[ - Union[BedrockGuardrailResponse, str] - ] = None + bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( + None + ) try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index 68f9dfd7abc..66b80c10f18 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -35,7 +35,10 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs @@ -179,6 +182,7 @@ class CustomCodeGuardrail(CustomGuardrail): self._compile_error = f"Failed to compile custom code: {e}" raise CustomCodeCompilationError(self._compile_error) from e + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 8e992297e5d..63541a1e2f9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -23,7 +23,10 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -483,6 +486,7 @@ class EnkryptAIGuardrails(CustomGuardrail): request_data=data, guardrail_name=self.guardrail_name ) + @log_guardrail_information async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", 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 b37074e25e7..9018675d7a5 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 @@ -10,7 +10,10 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -150,6 +153,7 @@ class GenericGuardrailAPI(CustomGuardrail): return result_metadata + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 90f689ed23c..8955bffc125 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -9,7 +9,8 @@ from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, - ModifyResponseException + ModifyResponseException, + log_guardrail_information, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -108,7 +109,9 @@ class GraySwanGuardrail(CustomGuardrail): self.categories = categories self.policy_id = policy_id self.fail_open = True if fail_open is None else bool(fail_open) - self.guardrail_timeout = 30.0 if guardrail_timeout is None else float(guardrail_timeout) + self.guardrail_timeout = ( + 30.0 if guardrail_timeout is None else float(guardrail_timeout) + ) # Streaming configuration self.streaming_end_of_stream_only = streaming_end_of_stream_only @@ -155,6 +158,7 @@ class GraySwanGuardrail(CustomGuardrail): # Unified Guardrail Interface (works with ALL endpoints automatically) # ------------------------------------------------------------------ + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, @@ -208,7 +212,9 @@ class GraySwanGuardrail(CustomGuardrail): messages = [{"role": role, "content": text} for text in texts] # Get dynamic params from request metadata - dynamic_body = self.get_guardrail_dynamic_request_body_params(request_data) or {} + dynamic_body = ( + self.get_guardrail_dynamic_request_body_params(request_data) or {} + ) if dynamic_body: verbose_proxy_logger.debug( "Gray Swan Guardrail: dynamic extra_body=%s", safe_dumps(dynamic_body) @@ -271,12 +277,12 @@ class GraySwanGuardrail(CustomGuardrail): async def run_grayswan_guardrail(self, payload: dict) -> Dict[str, Any]: """ Run the GraySwan guardrail on a payload. - + This is a legacy method for testing purposes. - + Args: payload: The payload to scan - + Returns: Dict containing the GraySwan API response """ @@ -293,11 +299,11 @@ class GraySwanGuardrail(CustomGuardrail): ) -> None: """ Legacy method for processing GraySwan API responses. - + This method is maintained for backward compatibility with existing tests. It handles the test scenarios where responses need to be processed with knowledge of the request context (pre/during/post call hooks). - + Args: response_json: Response from GraySwan API data: Optional request data (for passthrough exceptions) @@ -365,7 +371,10 @@ class GraySwanGuardrail(CustomGuardrail): ) # If hook_type is provided and in pre/during call, raise exception - if hook_type in [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call]: + if hook_type in [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + ]: # Raise ModifyResponseException to short-circuit LLM call if data is None: data = {} @@ -540,7 +549,9 @@ class GraySwanGuardrail(CustomGuardrail): if isinstance(litellm_metadata, dict) and litellm_metadata: cleaned_litellm_metadata = dict(litellm_metadata) # cleaned_litellm_metadata.pop("user_api_key_auth", None) - sanitized = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + sanitized = safe_json_loads( + safe_dumps(cleaned_litellm_metadata), default={} + ) if isinstance(sanitized, dict) and sanitized: payload["litellm_metadata"] = sanitized @@ -566,7 +577,9 @@ class GraySwanGuardrail(CustomGuardrail): detection_info = detection_info[0] # Extract fields from detection_info dict - detection_dict: dict = detection_info if isinstance(detection_info, dict) else {} + detection_dict: dict = ( + detection_info if isinstance(detection_info, dict) else {} + ) violation_score = detection_dict.get("violation_score", 0.0) violated_rules = detection_dict.get("violated_rules", []) mutation = detection_dict.get("mutation", False) @@ -582,7 +595,9 @@ class GraySwanGuardrail(CustomGuardrail): if violated_rules: formatted_rules = self._format_violated_rules(violated_rules) if formatted_rules: - message_parts.append(f"It was violating the rule(s): {formatted_rules}.") + message_parts.append( + f"It was violating the rule(s): {formatted_rules}." + ) if mutation: message_parts.append( @@ -590,9 +605,7 @@ class GraySwanGuardrail(CustomGuardrail): ) if ipi: - message_parts.append( - "Indirect Prompt Injection was DETECTED." - ) + message_parts.append("Indirect Prompt Injection was DETECTED.") return "\n".join(message_parts) diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index e2c20604880..b907fbbcbda 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -10,7 +10,10 @@ from httpx import HTTPStatusError from requests.auth import HTTPBasicAuth from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -110,6 +113,7 @@ class HiddenlayerGuardrail(CustomGuardrail): ) super().__init__(**kwargs) + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, 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 083a407e9cf..263b6eee768 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 @@ -28,7 +28,10 @@ from fastapi import HTTPException from litellm import Router from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ModelResponseStream @@ -50,6 +53,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor ContentFilterDetection, PatternDetection, ) + from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern MAX_KEYWORD_VALUE_GAP_WORDS = 1 @@ -168,9 +172,9 @@ class ContentFilterGuardrail(CustomGuardrail): self.image_model = image_model # Store loaded categories self.loaded_categories: Dict[str, CategoryConfig] = {} - self.category_keywords: Dict[ - str, Tuple[str, str, ContentFilterAction] - ] = {} # keyword -> (category, severity, action) + self.category_keywords: Dict[str, Tuple[str, str, ContentFilterAction]] = ( + {} + ) # keyword -> (category, severity, action) # Load categories if provided if categories: @@ -994,6 +998,7 @@ class ContentFilterGuardrail(CustomGuardrail): masked_entity_count=masked_entity_count, ) + @log_guardrail_information async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 3598dbe741e..1cfc805dbf9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -12,7 +12,10 @@ import httpx from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -26,7 +29,11 @@ if TYPE_CHECKING: class OnyxGuardrail(CustomGuardrail): def __init__( - self, api_base: Optional[str] = None, api_key: Optional[str] = None, timeout: Optional[float] = 10.0, **kwargs + self, + api_base: Optional[str] = None, + api_key: Optional[str] = None, + timeout: Optional[float] = 10.0, + **kwargs, ): timeout = timeout or int(os.getenv("ONYX_TIMEOUT", 10.0)) self.async_handler = get_async_httpx_client( @@ -79,6 +86,7 @@ class OnyxGuardrail(CustomGuardrail): ) return result + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 030b6036815..a196937ef6c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -58,7 +58,9 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): guardrail_name: str, api_key: Optional[str] = None, api_base: Optional[str] = None, - model: Optional[Literal["omni-moderation-latest", "text-moderation-latest"]] = None, + model: Optional[ + Literal["omni-moderation-latest", "text-moderation-latest"] + ] = None, **kwargs, ): """Initialize OpenAI Moderation guardrail handler.""" @@ -75,7 +77,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): supported_event_hooks=supported_event_hooks, **kwargs, ) - + self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback ) @@ -83,10 +85,14 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): # Store configuration self.api_key = api_key or self._get_api_key() self.api_base = api_base or "https://api.openai.com/v1" - self.model: Literal["omni-moderation-latest", "text-moderation-latest"] = model or "omni-moderation-latest" + self.model: Literal["omni-moderation-latest", "text-moderation-latest"] = ( + model or "omni-moderation-latest" + ) if not self.api_key: - raise ValueError("OpenAI Moderation: api_key is required. Set OPENAI_API_KEY environment variable or pass it in configuration.") + raise ValueError( + "OpenAI Moderation: api_key is required. Set OPENAI_API_KEY environment variable or pass it in configuration." + ) verbose_proxy_logger.debug( f"Initialized OpenAI Moderation Guardrail: {guardrail_name} with model: {self.model}" @@ -98,7 +104,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): import litellm from litellm.secret_managers.main import get_secret_str - + return ( os.environ.get("OPENAI_API_KEY") or litellm.api_key @@ -106,21 +112,14 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): or get_secret_str("OPENAI_API_KEY") ) - async def async_make_request( - self, input_text: str - ) -> "OpenAIModerationResponse": + async def async_make_request(self, input_text: str) -> "OpenAIModerationResponse": """ Make a request to the OpenAI Moderation API. """ - request_body = { - "model": self.model, - "input": input_text - } - - verbose_proxy_logger.debug( - "OpenAI Moderation guard request: %s", request_body - ) - + request_body = {"model": self.model, "input": input_text} + + verbose_proxy_logger.debug("OpenAI Moderation guard request: %s", request_body) + response = await self.async_handler.post( url=f"{self.api_base}/moderations", headers={ @@ -133,7 +132,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): verbose_proxy_logger.debug( "OpenAI Moderation guard response: %s", response.json() ) - + if response.status_code != 200: raise HTTPException( status_code=response.status_code, @@ -144,9 +143,12 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): ) from litellm.types.llms.openai import OpenAIModerationResponse + return OpenAIModerationResponse(**response.json()) - def _check_moderation_result(self, moderation_response: "OpenAIModerationResponse") -> None: + def _check_moderation_result( + self, moderation_response: "OpenAIModerationResponse" + ) -> None: """ Check if the moderation response indicates harmful content and raise exception if needed. """ @@ -168,10 +170,10 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): } verbose_proxy_logger.warning( - "OpenAI Moderation: Content flagged for violations: %s", - violation_details + "OpenAI Moderation: Content flagged for violations: %s", + violation_details, ) - + raise HTTPException( status_code=400, detail={ @@ -180,6 +182,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): }, ) + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, @@ -189,51 +192,50 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): ) -> GenericGuardrailAPIInputs: """ Apply OpenAI moderation guardrail using the unified guardrail interface. - + This method is called by the UnifiedLLMGuardrails system for all endpoint types (chat completions, embeddings, responses API, etc.). - + Args: inputs: GenericGuardrailAPIInputs containing texts and/or structured_messages request_data: The original request data input_type: Whether this is a "request" (pre-call) or "response" (post-call) logging_obj: Optional logging object - + Returns: The inputs unchanged (moderation doesn't modify content, only blocks) - + Raises: HTTPException: If content violates moderation policy """ # Extract text to moderate from inputs text_to_moderate: Optional[str] = None - + # Prefer structured_messages if available (has role context) if structured_messages := inputs.get("structured_messages"): text_to_moderate = self.get_user_prompt(structured_messages) - + # Fall back to texts if not text_to_moderate: if texts := inputs.get("texts"): # Join all texts for moderation text_to_moderate = "\n".join(texts) - + if not text_to_moderate: verbose_proxy_logger.debug( "OpenAI Moderation: No text content to moderate in inputs" ) return inputs - + # Make moderation request moderation_response = await self.async_make_request(input_text=text_to_moderate) - + # Check if content is flagged and raise exception if needed self._check_moderation_result(moderation_response) - + # Moderation doesn't modify content, just blocks - return inputs unchanged return inputs - @log_guardrail_information async def async_post_call_streaming_iterator_hook( self, @@ -252,9 +254,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): from litellm.main import stream_chunk_builder from litellm.types.utils import TextCompletionResponse - verbose_proxy_logger.debug( - "OpenAI Moderation: Running streaming response scan" - ) + verbose_proxy_logger.debug("OpenAI Moderation: Running streaming response scan") # Collect all chunks to process them together all_chunks: List["ModelResponseStream"] = [] @@ -269,7 +269,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): ) if isinstance(assembled_model_response, (type(None), TextCompletionResponse)): - # If we can't assemble a ModelResponse or it's a text completion, + # If we can't assemble a ModelResponse or it's a text completion, # just yield the original chunks without moderation verbose_proxy_logger.warning( "OpenAI Moderation: Could not assemble ModelResponse from chunks, skipping moderation" @@ -284,19 +284,17 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): verbose_proxy_logger.debug( f"OpenAI Moderation: Streaming response text: {response_text[:100]}..." # Log first 100 chars ) - + # Make moderation request - this will raise HTTPException if content is flagged moderation_response = await self.async_make_request( input_text=response_text, ) - + # Check if content is flagged and raise exception if needed self._check_moderation_result(moderation_response) # If we reach here, content passed moderation - yield the original chunks - mock_response = MockResponseIterator( - model_response=assembled_model_response - ) + mock_response = MockResponseIterator(model_response=assembled_model_response) # Return the reconstructed stream async for chunk in mock_response: @@ -306,34 +304,34 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): """ Extract text content from the model response for moderation. """ - if not hasattr(response, 'choices') or not response.choices: + if not hasattr(response, "choices") or not response.choices: return None response_texts = [] for choice in response.choices: try: # Try to get content from message (chat completion) - message = getattr(choice, 'message', None) + message = getattr(choice, "message", None) if message: - content = getattr(message, 'content', None) + content = getattr(message, "content", None) if content and isinstance(content, str): response_texts.append(content) continue - + # Try to get text (text completion) - text = getattr(choice, 'text', None) + text = getattr(choice, "text", None) if text and isinstance(text, str): response_texts.append(text) continue - + # Try to get content from delta (streaming) - delta = getattr(choice, 'delta', None) + delta = getattr(choice, "delta", None) if delta: - content = getattr(delta, 'content', None) + content = getattr(delta, "content", None) if content and isinstance(content, str): response_texts.append(content) continue - + except (AttributeError, TypeError): # Skip choices that don't have expected attributes continue diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 71ad9819146..3984384aae4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -9,10 +9,10 @@ import asyncio -import threading import json -from datetime import datetime +import threading from contextlib import asynccontextmanager +from datetime import datetime from typing import ( TYPE_CHECKING, Any, @@ -38,8 +38,11 @@ if TYPE_CHECKING: from litellm._uuid import uuid from litellm.caching.caching import DualCache -from litellm.exceptions import BlockedPiiEntityError -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import ( GuardrailEventHooks, @@ -229,6 +232,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """Cleanup: we try to close, but doing async cleanup in __del__ is risky.""" pass + def _has_block_action(self) -> bool: + """Return True if pii_entities_config has any BLOCK action (fail-closed on analyzer errors).""" + if not self.pii_entities_config: + return False + return any( + action == PiiAction.BLOCK for action in self.pii_entities_config.values() + ) + def _get_presidio_analyze_request_payload( self, text: str, @@ -313,13 +324,30 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # Handle error responses from Presidio (e.g., {'error': 'No text provided'}) # Presidio may return a dict instead of a list when errors occur + def _fail_on_invalid_response( + reason: str, + ) -> List[PresidioAnalyzeResponseItem]: + should_fail_closed = ( + bool(self.pii_entities_config) + or self.output_parse_pii + or self.apply_to_output + ) + if should_fail_closed: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Presidio analyzer returned invalid response; cannot verify PII when PII protection is configured: {reason}", + should_wrap_with_default_message=False, + ) + verbose_proxy_logger.warning( + "Presidio analyzer %s, returning empty list", reason + ) + return [] + if isinstance(analyze_results, dict): if "error" in analyze_results: - verbose_proxy_logger.warning( - "Presidio analyzer returned error: %s, returning empty list", - analyze_results.get("error"), + return _fail_on_invalid_response( + f"error: {analyze_results.get('error')}" ) - return [] # If it's a dict but not an error, try to process it as a single item verbose_proxy_logger.debug( "Presidio returned dict (not list), attempting to process as single item" @@ -327,23 +355,33 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): try: return [PresidioAnalyzeResponseItem(**analyze_results)] except Exception as e: - verbose_proxy_logger.warning( - "Failed to parse Presidio dict response: %s, returning empty list", - e, + return _fail_on_invalid_response( + f"failed to parse dict response: {e}" ) - return [] + + # Handle unexpected types (str, None, etc.) - e.g. from malformed/error + if not isinstance(analyze_results, list): + return _fail_on_invalid_response( + f"unexpected type {type(analyze_results).__name__} (expected list or dict), response: {str(analyze_results)[:200]}" + ) # Normal case: list of results final_results = [] for item in analyze_results: + if not isinstance(item, dict): + verbose_proxy_logger.warning( + "Skipping invalid Presidio result item (expected dict, got %s): %s", + type(item).__name__, + str(item)[:100], + ) + continue try: final_results.append(PresidioAnalyzeResponseItem(**item)) - except TypeError as te: - # Handle case where item is not a dict (shouldn't happen, but be defensive) + except Exception as e: verbose_proxy_logger.warning( - "Skipping invalid Presidio result item: %s (error: %s)", + "Failed to parse Presidio result item: %s (error: %s)", item, - te, + e, ) continue return final_results @@ -568,9 +606,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if messages is None: return data tasks = [] - task_mappings: List[ - Tuple[int, Optional[int]] - ] = [] # Track (message_index, content_index) for each task + task_mappings: List[Tuple[int, Optional[int]]] = ( + [] + ) # Track (message_index, content_index) for each task for msg_idx, m in enumerate(messages): content = m.get("content", None) @@ -671,9 +709,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ): # /chat/completions requests messages: Optional[List] = kwargs.get("messages", None) tasks = [] - task_mappings: List[ - Tuple[int, Optional[int]] - ] = [] # Track (message_index, content_index) for each task + task_mappings: List[Tuple[int, Optional[int]]] = ( + [] + ) # Track (message_index, content_index) for each task if messages is None: return kwargs, result @@ -792,11 +830,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # Type narrowing: StreamingChoices doesn't have .message attribute if not hasattr(choice, "message"): continue - content = getattr(choice.message, "content", None) + content = getattr(choice.message, "content", None) # type: ignore if content is None: continue if isinstance(content, str): - choice.message.content = await self.check_pii( + choice.message.content = await self.check_pii( # type: ignore text=content, output_parse_pii=False, presidio_config=presidio_config, @@ -989,6 +1027,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except Exception: pass + @log_guardrail_information async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 5ebc7b96eb8..b3e761869b0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -6,7 +6,10 @@ from typing import TYPE_CHECKING, Any, List, Literal, Optional, Type from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -67,6 +70,7 @@ class PromptSecurityGuardrail(CustomGuardrail): super().__init__(**kwargs) + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index 87da11efad0..6486da7f714 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -12,10 +12,11 @@ from typing import Any, Dict, List, Literal, Optional, Type from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.litellm_logging import ( - Logging as LiteLLMLoggingObj, +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, ) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -343,9 +344,7 @@ class QualifireGuardrail(CustomGuardrail): ) url = f"{self.qualifire_api_base}/api/evaluation/evaluate" - verbose_proxy_logger.debug( - f"Qualifire Guardrail: Making request to {url}" - ) + verbose_proxy_logger.debug(f"Qualifire Guardrail: Making request to {url}") # Make the API request response = await self.async_handler.post( @@ -393,6 +392,7 @@ class QualifireGuardrail(CustomGuardrail): verbose_proxy_logger.exception(f"Qualifire Guardrail error: {e}") raise + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index c60752d7952..ff00cd73ca5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -9,7 +9,10 @@ from typing import TYPE_CHECKING, Literal, Optional from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -70,6 +73,7 @@ class ZscalerAIGuard(CustomGuardrail): return str(value).strip() return "N/A" + @log_guardrail_information async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", @@ -92,7 +96,7 @@ class ZscalerAIGuard(CustomGuardrail): Raises: Exception: If content is blocked by Zscaler AI Guard """ - + texts = inputs.get("texts", []) try: verbose_proxy_logger.debug(f"ZscalerAIGuard: Checking {len(texts)} text(s)") @@ -102,8 +106,8 @@ class ZscalerAIGuard(CustomGuardrail): team_metadata = metadata.get("team_metadata", {}) or {} # Precedence for policy_id: - # 1. metadata.zguard_policy_id # request level - # 2. user_api_key_metadata.zguard_policy_id # Key level + # 1. metadata.zguard_policy_id # request level + # 2. user_api_key_metadata.zguard_policy_id # Key level # 3. team_metadata.zguard_policy_id # Team level # 4. self.policy_id (from environment) # Global policy_id = ( @@ -154,9 +158,7 @@ class ZscalerAIGuard(CustomGuardrail): zscaler_ai_guard_result and zscaler_ai_guard_result.get("action") == "BLOCK" ): - blocking_info = zscaler_ai_guard_result.get( - "zscaler_ai_guard_response" - ) + blocking_info = zscaler_ai_guard_result.get("zscaler_ai_guard_response") error_message = f"Content blocked by Zscaler AI Guard: {self.extract_blocking_info(blocking_info)}" raise Exception(error_message) except Exception as e: diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index ecad8bc1b11..78a371ad66e 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,7 +17,7 @@ Quick summary: - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union from fastapi import HTTPException from pydantic import BaseModel @@ -241,6 +241,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): self, file_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> BatchFileUsage: """ Count number of requests and tokens in a batch input file. @@ -248,6 +249,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): Args: file_id: The file ID to read custom_llm_provider: The custom LLM provider to use for token encoding + user_api_key_dict: User authentication information for file access (required for managed files) Returns: BatchFileUsage with total_tokens and request_count @@ -257,6 +259,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_content = await litellm.afile_content( file_id=file_id, custom_llm_provider=custom_llm_provider, + user_api_key_dict=user_api_key_dict, ) file_content_as_dict = _get_file_content_as_dictionary( @@ -336,6 +339,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage = await self.count_input_file_usage( file_id=input_file_id, custom_llm_provider=custom_llm_provider, + user_api_key_dict=user_api_key_dict, ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 69c7e92d82e..b8c073dd061 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -153,7 +153,10 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): "standard_logging_object", None ) if standard_logging_payload is None: - raise ValueError("standard_logging_payload is required") + verbose_proxy_logger.debug( + "Skipping _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: standard_logging_payload is None" + ) + return _litellm_params: dict = kwargs.get("litellm_params", {}) or {} _metadata: dict = _litellm_params.get("metadata", {}) or {} diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 4a2c05f8590..4a8eb8e7419 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -144,6 +144,12 @@ async def image_generation( litellm_call_id=data.get("litellm_call_id", ""), status="success" ) ) + + ### CALL HOOKS ### - modify outgoing data (guardrails, otel, etc.) + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + ### RESPONSE HEADERS ### hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = hidden_params.get("model_id", None) or "" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9be78264e85..49d31c1efec 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1539,8 +1539,15 @@ def add_guardrails_from_policy_engine( """ from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.callback_utils import ( + add_policy_sources_to_metadata, add_policy_to_applied_policies_header, ) + from litellm.proxy.common_utils.http_parsing_utils import ( + get_tags_from_request_body, + ) + from litellm.proxy.policy_engine.attachment_registry import ( + get_attachment_registry, + ) from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver @@ -1561,20 +1568,31 @@ def add_guardrails_from_policy_engine( ) return - # Build context from request + # Extract tags using the shared helper (handles metadata / litellm_metadata, + # top-level tags, deduplication, and type filtering). + + all_tags = get_tags_from_request_body(data) or None + context = PolicyMatchContext( team_alias=user_api_key_dict.team_alias, key_alias=user_api_key_dict.key_alias, model=data.get("model"), + tags=all_tags, ) verbose_proxy_logger.debug( f"Policy engine: matching policies for context team_alias={context.team_alias}, " - f"key_alias={context.key_alias}, model={context.model}" + f"key_alias={context.key_alias}, model={context.model}, tags={context.tags}" ) - # Get matching policies via attachments - matching_policy_names = PolicyMatcher.get_matching_policies(context=context) + # Get matching policies via attachments (with match reasons for attribution) + attachment_registry = get_attachment_registry() + matches_with_reasons = attachment_registry.get_attached_policies_with_reasons( + context + ) + matching_policy_names = [m["policy_name"] for m in matches_with_reasons] + # Build reasons map: {"hipaa-policy": "tag:healthcare", ...} + policy_reasons = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons} verbose_proxy_logger.debug( f"Policy engine: matched policies via attachments: {matching_policy_names}" @@ -1607,6 +1625,16 @@ def add_guardrails_from_policy_engine( request_data=data, policy_name=policy_name ) + # Track policy attribution sources for x-litellm-policy-sources header + applied_reasons = { + name: policy_reasons[name] + for name in applied_policy_names + if name in policy_reasons + } + add_policy_sources_to_metadata( + request_data=data, policy_sources=applied_reasons + ) + # Resolve guardrails from matching policies resolved_guardrails = PolicyResolver.resolve_guardrails_for_context(context=context) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 24a41a2361b..942758e3bab 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -8,6 +8,7 @@ from litellm.proxy._types import ( LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_OrganizationTable, LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth, ) @@ -108,6 +109,154 @@ async def _user_has_admin_privileges( return False +def _org_admin_can_invite_user( + admin_user_obj: LiteLLM_UserTable, + target_user_obj: LiteLLM_UserTable, +) -> bool: + """ + Check if an org admin can invite the target user. + Target user must be in at least one org where the admin has org admin role. + + Args: + admin_user_obj: The admin user's full object (from get_user_object) + target_user_obj: The target user's full object (from get_user_object) + + Returns: + True if target user is in an org where admin has org admin role + """ + if admin_user_obj.organization_memberships is None: + return False + admin_org_ids = { + m.organization_id + for m in admin_user_obj.organization_memberships + if m.user_role == LitellmUserRoles.ORG_ADMIN.value + } + if not admin_org_ids: + return False + if target_user_obj.organization_memberships is None: + return False + target_org_ids = { + m.organization_id for m in target_user_obj.organization_memberships + } + return bool(admin_org_ids & target_org_ids) + + +async def _team_admin_can_invite_user( + user_api_key_dict: UserAPIKeyAuth, + admin_user_obj: LiteLLM_UserTable, + target_user_obj: LiteLLM_UserTable, + prisma_client: "PrismaClient", +) -> bool: + """ + Check if a team admin can invite the target user. + Target user must be in at least one team where the admin has team admin role. + + Args: + user_api_key_dict: The admin user's API key auth object + admin_user_obj: The admin user's full object (from get_user_object) + target_user_obj: The target user's full object (from get_user_object) + prisma_client: Prisma client for database operations + + Returns: + True if target user is in a team where admin has team admin role + """ + if not admin_user_obj.teams or len(admin_user_obj.teams) == 0: + return False + if not target_user_obj.teams or len(target_user_obj.teams) == 0: + return False + + teams = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": admin_user_obj.teams}} + ) + admin_team_ids = [ + team.team_id + for team in teams + if _is_user_team_admin( + user_api_key_dict=user_api_key_dict, + team_obj=LiteLLM_TeamTable(**team.model_dump()), + ) + ] + if not admin_team_ids: + return False + target_team_ids = set(target_user_obj.teams) + return bool(set(admin_team_ids) & target_team_ids) + + +async def admin_can_invite_user( + target_user_id: str, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: Optional["PrismaClient"] = None, + user_api_key_cache: Optional["DualCache"] = None, + proxy_logging_obj: Optional["ProxyLogging"] = None, +) -> bool: + """ + Check if the admin can create an invitation for the target user. + - Proxy admins: can invite any user + - Org admins: can only invite users in their org(s) + - Team admins: can only invite users in their team(s) + + Uses get_user_object for caching of both admin and target user objects. + + Args: + target_user_id: The user_id of the user to invite + user_api_key_dict: The admin user's API key auth object + prisma_client: Prisma client for database operations + user_api_key_cache: Cache for user API keys + proxy_logging_obj: Proxy logging object + + Returns: + True if user can invite the target user + """ + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return True + + if prisma_client is None or user_api_key_dict.user_id is None: + return False + + from litellm.caching import DualCache as DualCacheImport + from litellm.proxy.auth.auth_checks import get_user_object + + try: + cache = user_api_key_cache or DualCacheImport() + admin_user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + if admin_user_obj is None: + return False + + target_user_obj = await get_user_object( + user_id=target_user_id, + prisma_client=prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + if target_user_obj is None: + return False + + if _org_admin_can_invite_user(admin_user_obj, target_user_obj): + return True + + if await _team_admin_can_invite_user( + user_api_key_dict=user_api_key_dict, + admin_user_obj=admin_user_obj, + target_user_obj=target_user_obj, + prisma_client=prisma_client, + ): + return True + + return False + except Exception as e: + verbose_proxy_logger.debug( + f"Error checking invite permission for user {user_api_key_dict.user_id}: {e}" + ) + return False + + def _set_object_metadata_field( object_data: Union[ LiteLLM_TeamTable, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 90c2f7fdf62..597521ae773 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -262,6 +262,67 @@ if MCP_AVAILABLE: ) -> List[LiteLLM_MCPServerTable]: return [_redact_mcp_credentials(server) for server in mcp_servers] + def _is_restricted_virtual_key_request(user_api_key_dict: UserAPIKeyAuth) -> bool: + """Best-effort detection for route-restricted virtual keys. + + We treat a requestor as a "restricted" virtual key if `allowed_routes` + is a non-empty list. This matches the auth gate that blocks routes with + the error: "Virtual key is not allowed to call this route...". + """ + + allowed_routes = getattr(user_api_key_dict, "allowed_routes", None) + return isinstance(allowed_routes, list) and len(allowed_routes) > 0 + + def _sanitize_mcp_server_for_virtual_key( + mcp_server: LiteLLM_MCPServerTable, + ) -> LiteLLM_MCPServerTable: + """Return a minimally sufficient MCP server view for virtual keys. + + Security model: + - Virtual keys should be able to *discover* accessible servers. + - They should NOT receive sensitive configuration details like upstream + URLs, env vars, headers, commands/args, access-group names, or + credentials. + """ + + sanitized = _redact_mcp_credentials(mcp_server) + + # Remove potentially sensitive config + identity fields. + sanitized.url = None + sanitized.static_headers = None + sanitized.env = {} + sanitized.command = None + sanitized.args = [] + sanitized.extra_headers = [] + sanitized.allowed_tools = [] + sanitized.mcp_access_groups = [] + sanitized.teams = [] + + sanitized.authorization_url = None + sanitized.token_url = None + sanitized.registration_url = None + + sanitized.health_check_error = None + sanitized.last_health_check = None + + sanitized.created_by = None + sanitized.updated_by = None + sanitized.created_at = None + sanitized.updated_at = None + + # `mcp_info` is arbitrary metadata; keep only an explicit safe subset. + is_public = False + if isinstance(sanitized.mcp_info, dict): + is_public = bool(sanitized.mcp_info.get("is_public")) + sanitized.mcp_info = {"is_public": True} if is_public else None + + return sanitized + + def _sanitize_mcp_server_list_for_virtual_key( + mcp_servers: Iterable[LiteLLM_MCPServerTable], + ) -> List[LiteLLM_MCPServerTable]: + return [_sanitize_mcp_server_for_virtual_key(server) for server in mcp_servers] + def _inherit_credentials_from_existing_server( payload: NewMCPServerRequest, ) -> NewMCPServerRequest: @@ -504,8 +565,11 @@ if MCP_AVAILABLE: """ user_mcp_management_mode = _get_user_mcp_management_mode() + is_restricted_virtual_key = _is_restricted_virtual_key_request( + user_api_key_dict + ) - if user_mcp_management_mode == "view_all": + if user_mcp_management_mode == "view_all" and not is_restricted_virtual_key: servers = await global_mcp_server_manager.get_all_mcp_servers_unfiltered() redacted_mcp_servers = _redact_mcp_credentials_list(servers) else: @@ -531,6 +595,11 @@ if MCP_AVAILABLE: if server.mcp_info is None: server.mcp_info = {} server.mcp_info["is_public"] = True + + # Virtual keys only get a sanitized discovery view. + if is_restricted_virtual_key: + return _sanitize_mcp_server_list_for_virtual_key(redacted_mcp_servers) + return redacted_mcp_servers @router.get( @@ -625,6 +694,34 @@ if MCP_AVAILABLE: detail={"error": f"MCP Server with id {server_id} not found"}, ) + # Implement authz restriction from requested user + is_admin_view = _user_has_admin_view(user_api_key_dict) + is_restricted_virtual_key = _is_restricted_virtual_key_request( + user_api_key_dict + ) + + if not is_admin_view: + # Perform authz check BEFORE any health check (avoid side-effects for + # unauthorized callers). + mcp_server_records = await get_all_mcp_servers_for_user( + prisma_client, user_api_key_dict + ) + exists = does_mcp_server_exist(mcp_server_records, server_id) + + if not exists: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"User does not have permission to view mcp server with id {server_id}. " + "You can only view mcp servers that you have access to." + ) + }, + ) + + # At this point caller is authorized to view the server. + await global_mcp_server_manager.add_server(mcp_server) + # Perform health check on the server using server manager try: health_result = await global_mcp_server_manager.health_check_server( @@ -644,26 +741,10 @@ if MCP_AVAILABLE: mcp_server.last_health_check = datetime.now() mcp_server.health_check_error = str(e) - # Implement authz restriction from requested user - if _user_has_admin_view(user_api_key_dict): - return _redact_mcp_credentials(mcp_server) - - # Perform authz check to filter the mcp servers user has access to - mcp_server_records = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) - exists = does_mcp_server_exist(mcp_server_records, server_id) - - if exists: - await global_mcp_server_manager.add_server(mcp_server) - return _redact_mcp_credentials(mcp_server) - else: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": f"User does not have permission to view mcp server with id {server_id}. You can only view mcp servers that you have access to." - }, - ) + redacted = _redact_mcp_credentials(mcp_server) + if is_restricted_virtual_key: + return _sanitize_mcp_server_for_virtual_key(redacted) + return redacted @router.post( "/server", diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index a35fc4a5f3f..295c2ad50b3 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -4,7 +4,7 @@ Types for the management endpoints Might include fastapi/proxy requirements.txt related imports """ -from typing import List, Optional, cast +from typing import Any, Dict, List, Optional, cast from fastapi_sso.sso.base import OpenID @@ -56,3 +56,4 @@ def get_litellm_user_role(role_str) -> Optional[LitellmUserRoles]: class CustomOpenID(OpenID): team_ids: List[str] user_role: Optional[LitellmUserRoles] = None + extra_fields: Optional[Dict[str, Any]] = None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 278f3bdaafd..7274b389a92 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -190,9 +190,15 @@ def process_sso_jwt_access_token( if access_token_str and result: import jwt - access_token_payload = jwt.decode( - access_token_str, options={"verify_signature": False} - ) + try: + access_token_payload = jwt.decode( + access_token_str, options={"verify_signature": False} + ) + except jwt.exceptions.DecodeError: + verbose_proxy_logger.debug( + "Access token is not a valid JWT (possibly an opaque token), skipping JWT-based extraction" + ) + return # Extract team IDs from access token if sso_jwt_handler is available if sso_jwt_handler: @@ -401,6 +407,8 @@ def generic_response_convertor( generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") + generic_user_extra_attributes = os.getenv("GENERIC_USER_EXTRA_ATTRIBUTES", None) + verbose_proxy_logger.debug( f" generic_user_id_attribute_name: {generic_user_id_attribute_name}\n generic_user_email_attribute_name: {generic_user_email_attribute_name}" ) @@ -473,6 +481,14 @@ def generic_response_convertor( f"Found valid LitellmUserRoles '{role.value}' from SSO attribute '{generic_user_role_attribute_name}'" ) + # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified + extra_fields: Optional[Dict[str, Any]] = None + if generic_user_extra_attributes: + extra_fields = {} + for attr_name in generic_user_extra_attributes.split(","): + attr_name = attr_name.strip() + extra_fields[attr_name] = get_nested_value(response, attr_name) + return CustomOpenID( id=get_nested_value(response, generic_user_id_attribute_name), display_name=get_nested_value( @@ -484,6 +500,7 @@ def generic_response_convertor( provider=get_nested_value(response, generic_provider_attribute_name), team_ids=all_teams, user_role=user_role, + extra_fields=extra_fields, ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3dab6ea14f8..81144ad9f31 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -642,10 +642,10 @@ def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: by finding the action in the endpoint and extracting everything between "model" and the action. Args: - endpoint: The endpoint path (e.g., "/model/aws/anthropic/model-name/invoke") + endpoint: The endpoint path (e.g., "/model/aws/anthropic/model-name/invoke" or "v2/model/model-name/invoke") Returns: - The extracted model name (e.g., "aws/anthropic/model-name") + The extracted model name (e.g., "aws/anthropic/model-name" or "model-name") Raises: ValueError: If model cannot be extracted from endpoint @@ -657,7 +657,34 @@ def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: # Format: model/application-inference-profile/{profile-id}/{action} return "/".join(endpoint_parts[1:3]) - # Format: model/{modelId}/{action} + # Format: model/{modelId}/{action} or v2/model/{modelId}/{action} + # Find the index of "model" in the endpoint parts + model_index = None + for idx, part in enumerate(endpoint_parts): + if part == "model": + model_index = idx + break + + # If "model" keyword not found, try to extract model from the endpoint + # by finding the action and taking everything before it + if model_index is None: + # Find the index of the action in the endpoint parts + action_index = None + for idx, part in enumerate(endpoint_parts): + if part in BEDROCK_ENDPOINT_ACTIONS: + action_index = idx + break + + if action_index is not None and action_index > 1: + # Join all parts before the action (excluding empty strings) + model_parts = [p for p in endpoint_parts[1:action_index] if p] + if model_parts: + return "/".join(model_parts) + + raise ValueError( + f"'model' keyword not found and unable to extract model from endpoint. Expected format: /model/{{modelId}}/{{action}}. Got: {endpoint}" + ) + # Find the index of the action in the endpoint parts action_index = None for idx, part in enumerate(endpoint_parts): @@ -665,13 +692,22 @@ def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: action_index = idx break - if action_index is not None and action_index > 1: - # Join all parts between "model" and the action - return "/".join(endpoint_parts[1:action_index]) + if action_index is not None and action_index > model_index + 1: + # Join all parts between "model" and the action (excluding "model" itself) + return "/".join(endpoint_parts[model_index + 1:action_index]) # Fallback to taking everything after "model" if no action found - return "/".join(endpoint_parts[1:]) + model_parts = [p for p in endpoint_parts[model_index + 1:] if p] + if model_parts: + return "/".join(model_parts) + raise ValueError( + f"No model ID found after 'model' keyword. Expected format: /model/{{modelId}}/{{action}}. Got: {endpoint}" + ) + + except ValueError: + # Re-raise ValueError as-is + raise except Exception as e: raise ValueError( f"Model missing from endpoint. Expected format: /model/{{modelId}}/{{action}}. Got: {endpoint}" diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 51a7c37717c..a7b60c8b185 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2201,6 +2201,41 @@ async def initialize_pass_through_endpoints( InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_key) +def _get_pass_through_endpoints_from_config() -> List[PassThroughGenericEndpoint]: + """ + Get pass-through endpoints defined in the config file. + These are read-only and cannot be edited via the UI. + Malformed endpoints are logged and skipped; they do not crash the function. + """ + from pydantic import ValidationError + + from litellm.proxy.proxy_server import config_passthrough_endpoints + + if config_passthrough_endpoints is None or len(config_passthrough_endpoints) == 0: + return [] + + returned_endpoints: List[PassThroughGenericEndpoint] = [] + for endpoint in config_passthrough_endpoints: + try: + if isinstance(endpoint, dict): + endpoint_dict = dict(endpoint) + endpoint_dict["is_from_config"] = True + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + elif isinstance(endpoint, PassThroughGenericEndpoint): + # Create a copy with is_from_config=True + endpoint_dict = endpoint.model_dump() + endpoint_dict["is_from_config"] = True + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + except ValidationError as e: + verbose_proxy_logger.warning( + "Skipping malformed pass-through endpoint from config: %s", + e, + exc_info=False, + ) + + return returned_endpoints + + async def _get_pass_through_endpoints_from_db( endpoint_id: Optional[str] = None, user_api_key_dict: Optional[UserAPIKeyAuth] = None, @@ -2223,17 +2258,27 @@ async def _get_pass_through_endpoints_from_db( returned_endpoints: List[PassThroughGenericEndpoint] = [] if endpoint_id is None: - # Return all endpoints + # Return all endpoints from DB, mark as not from config for endpoint in pass_through_endpoint_data: if isinstance(endpoint, dict): - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint)) + endpoint_dict = dict(endpoint) + endpoint_dict["is_from_config"] = False + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) elif isinstance(endpoint, PassThroughGenericEndpoint): - returned_endpoints.append(endpoint) + endpoint_dict = endpoint.model_dump() + endpoint_dict["is_from_config"] = False + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) else: # Find specific endpoint by ID found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) if found_endpoint is not None: - returned_endpoints.append(found_endpoint) + endpoint_dict = ( + found_endpoint.model_dump() + if isinstance(found_endpoint, PassThroughGenericEndpoint) + else dict(found_endpoint) + ) + endpoint_dict["is_from_config"] = False + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) return returned_endpoints @@ -2312,10 +2357,25 @@ async def get_pass_through_endpoints( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - pass_through_endpoints = await _get_pass_through_endpoints_from_db( + # Get endpoints from DB (editable via UI) + db_endpoints = await _get_pass_through_endpoints_from_db( endpoint_id=endpoint_id, user_api_key_dict=user_api_key_dict ) + # Get endpoints from config file (read-only, not editable via UI) + config_endpoints = _get_pass_through_endpoints_from_config() + + # Merge: config endpoints not in DB + all DB endpoints (DB overrides config for same path) + db_paths = {ep.path for ep in db_endpoints} + config_only_endpoints = [ + ep for ep in config_endpoints if ep.path not in db_paths + ] + if endpoint_id is not None: + # When filtering by endpoint_id, only return if found in DB (config endpoints use generated IDs) + pass_through_endpoints = db_endpoints + else: + pass_through_endpoints = config_only_endpoints + db_endpoints + if team_id is not None: pass_through_endpoints = await _filter_endpoints_by_team_allowed_routes( team_id=team_id, @@ -2392,7 +2452,8 @@ async def update_pass_through_endpoints( ) # Get the update data as dict, excluding None values for partial updates - update_data = data.model_dump(exclude_none=True) + # Exclude is_from_config as it's a response-only field (computed at read time) + update_data = data.model_dump(exclude_none=True, exclude={"is_from_config"}) # Start with existing endpoint data endpoint_dict = found_endpoint.model_dump() @@ -2404,6 +2465,9 @@ async def update_pass_through_endpoints( if "id" not in update_data and found_endpoint.id is not None: endpoint_dict["id"] = found_endpoint.id + # Remove is_from_config before saving - it's a response-only field (computed at read time) + endpoint_dict.pop("is_from_config", None) + # Create updated endpoint object updated_endpoint = PassThroughGenericEndpoint(**endpoint_dict) @@ -2490,7 +2554,8 @@ async def create_pass_through_endpoints( ) ## Auto-generate ID if not provided - data_dict = data.model_dump() + # Exclude is_from_config as it's a response-only field (computed at read time) + data_dict = data.model_dump(exclude={"is_from_config"}) if data_dict.get("id") is None: data_dict["id"] = str(uuid.uuid4()) diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 4a335b54747..69b3b3599f3 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -84,6 +84,7 @@ class AttachmentRegistry: teams=attachment_data.get("teams"), keys=attachment_data.get("keys"), models=attachment_data.get("models"), + tags=attachment_data.get("tags"), ) def get_attached_policies(self, context: PolicyMatchContext) -> List[str]: @@ -96,21 +97,68 @@ class AttachmentRegistry: Returns: List of policy names that are attached to matching scopes """ + return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)] + + def get_attached_policies_with_reasons( + self, context: PolicyMatchContext + ) -> List[Dict[str, Any]]: + """ + Get list of policy names and match reasons for the given context. + + Returns a list of dicts with 'policy_name' and 'matched_via' keys. + The 'matched_via' describes which dimension caused the match. + """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - attached_policies: List[str] = [] + results: List[Dict[str, Any]] = [] + seen_policies: set = set() for attachment in self._attachments: scope = attachment.to_policy_scope() if PolicyMatcher.scope_matches(scope=scope, context=context): - if attachment.policy not in attached_policies: - attached_policies.append(attachment.policy) + if attachment.policy not in seen_policies: + seen_policies.add(attachment.policy) + matched_via = self._describe_match_reason(attachment, context) + results.append( + { + "policy_name": attachment.policy, + "matched_via": matched_via, + } + ) verbose_proxy_logger.debug( f"Attachment matched: policy={attachment.policy}, " + f"matched_via={matched_via}, " f"context=(team={context.team_alias}, key={context.key_alias}, model={context.model})" ) - return attached_policies + return results + + @staticmethod + def _describe_match_reason( + attachment: PolicyAttachment, context: PolicyMatchContext + ) -> str: + """Describe why an attachment matched the context.""" + from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher + + if attachment.is_global(): + return "scope:*" + + reasons = [] + if attachment.tags and context.tags: + matching_tags = [ + t for t in context.tags + if PolicyMatcher.matches_pattern(t, attachment.tags) + ] + if matching_tags: + reasons.append(f"tag:{matching_tags[0]}") + if attachment.teams and context.team_alias: + reasons.append(f"team:{context.team_alias}") + if attachment.keys and context.key_alias: + reasons.append(f"key:{context.key_alias}") + if attachment.models and context.model: + reasons.append(f"model:{context.model}") + + return "+".join(reasons) if reasons else "scope:default" def is_policy_attached( self, policy_name: str, context: PolicyMatchContext @@ -238,6 +286,7 @@ class AttachmentRegistry: "teams": attachment_request.teams or [], "keys": attachment_request.keys or [], "models": attachment_request.models or [], + "tags": attachment_request.tags or [], "created_at": datetime.now(timezone.utc), "updated_at": datetime.now(timezone.utc), "created_by": created_by, @@ -253,6 +302,7 @@ class AttachmentRegistry: teams=attachment_request.teams, keys=attachment_request.keys, models=attachment_request.models, + tags=attachment_request.tags, ) self.add_attachment(attachment) @@ -263,6 +313,7 @@ class AttachmentRegistry: teams=created_attachment.teams or [], keys=created_attachment.keys or [], models=created_attachment.models or [], + tags=created_attachment.tags or [], created_at=created_attachment.created_at, updated_at=created_attachment.updated_at, created_by=created_attachment.created_by, @@ -344,6 +395,7 @@ class AttachmentRegistry: teams=attachment.teams or [], keys=attachment.keys or [], models=attachment.models or [], + tags=attachment.tags or [], created_at=attachment.created_at, updated_at=attachment.updated_at, created_by=attachment.created_by, @@ -381,6 +433,7 @@ class AttachmentRegistry: teams=a.teams or [], keys=a.keys or [], models=a.models or [], + tags=a.tags or [], created_at=a.created_at, updated_at=a.updated_at, created_by=a.created_by, @@ -415,6 +468,7 @@ class AttachmentRegistry: teams=attachment_response.teams if attachment_response.teams else None, keys=attachment_response.keys if attachment_response.keys else None, models=attachment_response.models if attachment_response.models else None, + tags=attachment_response.tags if attachment_response.tags else None, ) self._attachments.append(attachment) diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index 615e153862a..3bd893b0034 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -23,10 +23,6 @@ from litellm.types.proxy.policy_engine import ( router = APIRouter() -# Get singleton instances -POLICY_REGISTRY = get_policy_registry() -ATTACHMENT_REGISTRY = get_attachment_registry() - # ───────────────────────────────────────────────────────────────────────────── # Policy CRUD Endpoints @@ -75,7 +71,7 @@ async def list_policies(): raise HTTPException(status_code=500, detail="Database not connected") try: - policies = await POLICY_REGISTRY.get_all_policies_from_db(prisma_client) + policies = await get_policy_registry().get_all_policies_from_db(prisma_client) return PolicyListDBResponse(policies=policies, total_count=len(policies)) except Exception as e: verbose_proxy_logger.exception(f"Error listing policies: {e}") @@ -130,7 +126,7 @@ async def create_policy( try: created_by = user_api_key_dict.user_id - result = await POLICY_REGISTRY.add_policy_to_db( + result = await get_policy_registry().add_policy_to_db( policy_request=request, prisma_client=prisma_client, created_by=created_by, @@ -168,7 +164,7 @@ async def get_policy(policy_id: str): raise HTTPException(status_code=500, detail="Database not connected") try: - result = await POLICY_REGISTRY.get_policy_by_id_from_db( + result = await get_policy_registry().get_policy_by_id_from_db( policy_id=policy_id, prisma_client=prisma_client, ) @@ -216,7 +212,7 @@ async def update_policy( try: # Check if policy exists - existing = await POLICY_REGISTRY.get_policy_by_id_from_db( + existing = await get_policy_registry().get_policy_by_id_from_db( policy_id=policy_id, prisma_client=prisma_client, ) @@ -226,7 +222,7 @@ async def update_policy( ) updated_by = user_api_key_dict.user_id - result = await POLICY_REGISTRY.update_policy_in_db( + result = await get_policy_registry().update_policy_in_db( policy_id=policy_id, policy_request=request, prisma_client=prisma_client, @@ -269,7 +265,7 @@ async def delete_policy(policy_id: str): try: # Check if policy exists - existing = await POLICY_REGISTRY.get_policy_by_id_from_db( + existing = await get_policy_registry().get_policy_by_id_from_db( policy_id=policy_id, prisma_client=prisma_client, ) @@ -278,7 +274,7 @@ async def delete_policy(policy_id: str): status_code=404, detail=f"Policy with ID {policy_id} not found" ) - result = await POLICY_REGISTRY.delete_policy_from_db( + result = await get_policy_registry().delete_policy_from_db( policy_id=policy_id, prisma_client=prisma_client, ) @@ -324,7 +320,7 @@ async def get_resolved_guardrails(policy_id: str): try: # Get the policy - policy = await POLICY_REGISTRY.get_policy_by_id_from_db( + policy = await get_policy_registry().get_policy_by_id_from_db( policy_id=policy_id, prisma_client=prisma_client, ) @@ -334,7 +330,7 @@ async def get_resolved_guardrails(policy_id: str): ) # Resolve guardrails - resolved = await POLICY_REGISTRY.resolve_guardrails_from_db( + resolved = await get_policy_registry().resolve_guardrails_from_db( policy_name=policy.policy_name, prisma_client=prisma_client, ) @@ -399,7 +395,7 @@ async def list_policy_attachments(): raise HTTPException(status_code=500, detail="Database not connected") try: - attachments = await ATTACHMENT_REGISTRY.get_all_attachments_from_db( + attachments = await get_attachment_registry().get_all_attachments_from_db( prisma_client ) return PolicyAttachmentListResponse( @@ -466,7 +462,7 @@ async def create_policy_attachment( try: # Verify the policy exists - policy = await POLICY_REGISTRY.get_all_policies_from_db(prisma_client) + policy = await get_policy_registry().get_all_policies_from_db(prisma_client) policy_names = [p.policy_name for p in policy] if request.policy_name not in policy_names: raise HTTPException( @@ -475,7 +471,7 @@ async def create_policy_attachment( ) created_by = user_api_key_dict.user_id - result = await ATTACHMENT_REGISTRY.add_attachment_to_db( + result = await get_attachment_registry().add_attachment_to_db( attachment_request=request, prisma_client=prisma_client, created_by=created_by, @@ -510,7 +506,7 @@ async def get_policy_attachment(attachment_id: str): raise HTTPException(status_code=500, detail="Database not connected") try: - result = await ATTACHMENT_REGISTRY.get_attachment_by_id_from_db( + result = await get_attachment_registry().get_attachment_by_id_from_db( attachment_id=attachment_id, prisma_client=prisma_client, ) @@ -556,7 +552,7 @@ async def delete_policy_attachment(attachment_id: str): try: # Check if attachment exists - existing = await ATTACHMENT_REGISTRY.get_attachment_by_id_from_db( + existing = await get_attachment_registry().get_attachment_by_id_from_db( attachment_id=attachment_id, prisma_client=prisma_client, ) @@ -566,7 +562,7 @@ async def delete_policy_attachment(attachment_id: str): detail=f"Attachment with ID {attachment_id} not found", ) - result = await ATTACHMENT_REGISTRY.delete_attachment_from_db( + result = await get_attachment_registry().delete_attachment_from_db( attachment_id=attachment_id, prisma_client=prisma_client, ) diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index ab73970bfab..888981f85f5 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -81,6 +81,19 @@ class PolicyMatcher: if not PolicyMatcher.matches_pattern(context.model, scope.get_models()): return False + # Check tags (only if scope specifies tags) + # Unlike teams/keys/models, empty tags means "do not check" rather than "match all" + scope_tags = scope.get_tags() + if scope_tags: + if not context.tags: + return False + # Match if ANY context tag matches ANY scope tag pattern + if not any( + PolicyMatcher.matches_pattern(tag, scope_tags) + for tag in context.tags + ): + return False + return True @staticmethod diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 5fb5084f648..a2431977b24 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -484,6 +484,7 @@ class PolicyRegistry: ) self.add_policy(policy_response.policy_name, policy) + self._initialized = True verbose_proxy_logger.info( f"Synced {len(policies)} policies from DB to in-memory registry" ) diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py new file mode 100644 index 00000000000..eb4d3fc5845 --- /dev/null +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -0,0 +1,408 @@ +""" +Policy resolve and attachment impact estimation endpoints. + +- /policies/resolve — debug which guardrails apply for a given context +- /policies/attachments/estimate-impact — preview blast radius before creating an attachment +""" + +import json + +from fastapi import APIRouter, Depends, HTTPException, Query + +from litellm._logging import verbose_proxy_logger +from litellm.constants import MAX_POLICY_ESTIMATE_IMPACT_ROWS +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.types.proxy.policy_engine import ( + AttachmentImpactResponse, + PolicyAttachmentCreateRequest, + PolicyMatchContext, + PolicyMatchDetail, + PolicyResolveRequest, + PolicyResolveResponse, +) + +router = APIRouter() + + +def _build_alias_where(field: str, patterns: list) -> dict: + """Build a Prisma ``where`` clause for alias patterns. + + Supports exact matches and suffix wildcards (``prefix*``). + Returns something like: + {"OR": [{"field": {"in": ["a","b"]}}, {"field": {"startsWith": "dev-"}}]} + """ + exact: list = [] + prefix_conditions: list = [] + for pat in patterns: + if pat.endswith("*"): + prefix_conditions.append({field: {"startsWith": pat[:-1]}}) + else: + exact.append(pat) + + conditions: list = [] + if exact: + conditions.append({field: {"in": exact}}) + conditions.extend(prefix_conditions) + + if not conditions: + return {field: {"not": None}} + if len(conditions) == 1: + return conditions[0] + return {"OR": conditions} + + +def _parse_metadata(raw_metadata: object) -> dict: + """Parse metadata that may be a dict, JSON string, or None.""" + if raw_metadata is None: + return {} + if isinstance(raw_metadata, str): + try: + return json.loads(raw_metadata) + except (json.JSONDecodeError, TypeError): + return {} + return raw_metadata if isinstance(raw_metadata, dict) else {} + + +def _get_tags_from_metadata(metadata: object, json_metadata: object = None) -> list: + """Extract tags list from a metadata field (or metadata_json fallback).""" + raw = json_metadata if json_metadata is not None else metadata + parsed = _parse_metadata(raw) + return parsed.get("tags", []) or [] + + +async def _fetch_all_teams(prisma_client: object) -> list: + """Fetch teams from DB once. Reuse the result across tag and alias lookups.""" + return await prisma_client.db.litellm_teamtable.find_many( # type: ignore + where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + ) + + +def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple: + """Filter key rows whose metadata.tags match any of the given patterns. + + Returns (named_aliases, unnamed_count). + """ + from litellm.proxy.auth.route_checks import RouteChecks + + affected: list = [] + unnamed_count = 0 + for key in keys: + key_alias = key.key_alias or "" + key_tags = _get_tags_from_metadata( + key.metadata, getattr(key, "metadata_json", None) + ) + if key_tags and any( + RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + for tag in key_tags + for pat in tag_patterns + ): + if key_alias: + affected.append(key_alias) + else: + unnamed_count += 1 + return affected, unnamed_count + + +def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple: + """Filter pre-fetched team rows whose metadata.tags match any patterns. + + Returns (named_aliases, unnamed_count). + """ + from litellm.proxy.auth.route_checks import RouteChecks + + affected: list = [] + unnamed_count = 0 + for team in teams: + team_alias = team.team_alias or "" + team_tags = _get_tags_from_metadata(team.metadata) + if team_tags and any( + RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + for tag in team_tags + for pat in tag_patterns + ): + if team_alias: + affected.append(team_alias) + else: + unnamed_count += 1 + return affected, unnamed_count + + +async def _find_affected_by_team_patterns( + prisma_client: object, + all_teams: list, + team_patterns: list, + existing_teams: list, + existing_keys: list, +) -> tuple: + """Filter pre-fetched teams by alias patterns, then fetch their keys. + + Returns (new_teams, new_keys, unnamed_keys_count). + """ + from litellm.proxy.auth.route_checks import RouteChecks + + new_teams: list = [] + matched_team_ids: list = [] + + for team in all_teams: + team_alias = team.team_alias or "" + if team_alias and any( + RouteChecks._route_matches_wildcard_pattern(route=team_alias, pattern=pat) + for pat in team_patterns + ): + if team_alias not in existing_teams: + new_teams.append(team_alias) + matched_team_ids.append(str(team.team_id)) + + new_keys: list = [] + unnamed_keys_count = 0 + if matched_team_ids: + keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + where={"team_id": {"in": matched_team_ids}}, + order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + ) + for key in keys: + key_alias = key.key_alias or "" + if key_alias: + if key_alias not in existing_keys: + new_keys.append(key_alias) + else: + unnamed_keys_count += 1 + + return new_teams, new_keys, unnamed_keys_count + + +async def _find_affected_keys_by_alias( + prisma_client: object, key_patterns: list, existing_keys: list +) -> list: + """Find keys whose alias matches the given patterns.""" + from litellm.proxy.auth.route_checks import RouteChecks + + affected: list = [] + + keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + where=_build_alias_where("key_alias", key_patterns), + order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + ) + for key in keys: + key_alias = key.key_alias or "" + if key_alias and any( + RouteChecks._route_matches_wildcard_pattern(route=key_alias, pattern=pat) + for pat in key_patterns + ): + if key_alias not in existing_keys: + affected.append(key_alias) + return affected + + +# ───────────────────────────────────────────────────────────────────────────── +# Policy Resolve Endpoint +# ───────────────────────────────────────────────────────────────────────────── + + +@router.post( + "/policies/resolve", + tags=["Policies"], + dependencies=[Depends(user_api_key_auth)], + response_model=PolicyResolveResponse, +) +async def resolve_policies_for_context( + request: PolicyResolveRequest, + force_sync: bool = Query( + default=False, + description="Force a DB sync before resolving. Default uses in-memory cache.", + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Resolve which policies and guardrails apply for a given context. + + Use this endpoint to debug "what guardrails would apply to a request + with this team/key/model/tags combination?" + + Example Request: + ```bash + curl -X POST "http://localhost:4000/policies/resolve" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "tags": ["healthcare"], + "model": "gpt-4" + }' + ``` + """ + from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher + from litellm.proxy.policy_engine.policy_resolver import PolicyResolver + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + # Only sync from DB when explicitly requested; otherwise use in-memory cache + if force_sync: + await get_policy_registry().sync_policies_from_db(prisma_client) + await get_attachment_registry().sync_attachments_from_db(prisma_client) + + # Build context from request + context = PolicyMatchContext( + team_alias=request.team_alias, + key_alias=request.key_alias, + model=request.model, + tags=request.tags, + ) + + # Get matching policies with reasons + match_results = get_attachment_registry().get_attached_policies_with_reasons( + context=context + ) + + if not match_results: + return PolicyResolveResponse( + effective_guardrails=[], + matched_policies=[], + ) + + # Filter by conditions + policy_names = [r["policy_name"] for r in match_results] + applied_policy_names = PolicyMatcher.get_policies_with_matching_conditions( + policy_names=policy_names, + context=context, + ) + + # Resolve guardrails for each applied policy + matched_policies = [] + all_guardrails: set = set() + for result in match_results: + pname = result["policy_name"] + if pname not in applied_policy_names: + continue + resolved = PolicyResolver.resolve_policy_guardrails( + policy_name=pname, + policies=get_policy_registry().get_all_policies(), + context=context, + ) + guardrails = resolved.guardrails if resolved else [] + all_guardrails.update(guardrails) + matched_policies.append( + PolicyMatchDetail( + policy_name=pname, + matched_via=result["matched_via"], + guardrails_added=guardrails, + ) + ) + + return PolicyResolveResponse( + effective_guardrails=sorted(all_guardrails), + matched_policies=matched_policies, + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception(f"Error resolving policies: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ───────────────────────────────────────────────────────────────────────────── +# Attachment Impact Estimation Endpoint +# ───────────────────────────────────────────────────────────────────────────── + + +@router.post( + "/policies/attachments/estimate-impact", + tags=["Policies"], + dependencies=[Depends(user_api_key_auth)], + response_model=AttachmentImpactResponse, +) +async def estimate_attachment_impact( + request: PolicyAttachmentCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Estimate how many keys and teams would be affected by a policy attachment. + + Use this before creating an attachment to preview the blast radius. + + Example Request: + ```bash + curl -X POST "http://localhost:4000/policies/attachments/estimate-impact" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "policy_name": "hipaa-compliance", + "tags": ["healthcare", "health-*"] + }' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + # If global scope, everything is affected — not useful to enumerate + if request.scope == "*": + return AttachmentImpactResponse( + affected_keys_count=-1, + affected_teams_count=-1, + sample_keys=["(global scope — affects all keys)"], + sample_teams=["(global scope — affects all teams)"], + ) + + affected_keys: list = [] + affected_teams: list = [] + unnamed_keys = 0 + unnamed_teams = 0 + + tag_patterns = request.tags or [] + team_patterns = request.teams or [] + + # Fetch teams once — reused by both tag-based and alias-based lookups + all_teams: list = [] + if tag_patterns or team_patterns: + all_teams = await _fetch_all_teams(prisma_client) + + # Tag-based impact + if tag_patterns: + keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + where={}, order={"created_at": "desc"}, + take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + ) + affected_keys, unnamed_keys = _filter_keys_by_tags(keys, tag_patterns) + affected_teams, unnamed_teams = _filter_teams_by_tags( + all_teams, tag_patterns, + ) + + # Team-based impact (alias matching + keys belonging to those teams) + if team_patterns: + new_teams, new_keys, new_unnamed = await _find_affected_by_team_patterns( + prisma_client, all_teams, team_patterns, + affected_teams, affected_keys, + ) + affected_teams.extend(new_teams) + affected_keys.extend(new_keys) + unnamed_keys += new_unnamed + + # Key-based impact (direct alias matching) + key_patterns = request.keys or [] + if key_patterns: + new_keys = await _find_affected_keys_by_alias( + prisma_client, key_patterns, affected_keys, + ) + affected_keys.extend(new_keys) + + return AttachmentImpactResponse( + affected_keys_count=len(affected_keys) + unnamed_keys, + affected_teams_count=len(affected_teams) + unnamed_teams, + unnamed_keys_count=unnamed_keys, + unnamed_teams_count=unnamed_teams, + sample_keys=affected_keys[:10], + sample_teams=affected_teams[:10], + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception(f"Error estimating attachment impact: {e}") + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index d87ae8b14ca..a094eb84bf3 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -46,6 +46,7 @@ mcp_servers: transport: "http" url: "https://mcp.deepwiki.com/mcp" + # General Settings general_settings: master_key: sk-1234 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 294294cdda7..6286d6dd1ca 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -338,7 +338,10 @@ from litellm.proxy.management_endpoints.cache_settings_endpoints import ( from litellm.proxy.management_endpoints.callback_management_endpoints import ( router as callback_management_endpoints_router, ) -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.proxy.management_endpoints.common_utils import ( + admin_can_invite_user, + _user_has_admin_privileges, +) from litellm.proxy.management_endpoints.cost_tracking_settings import ( router as cost_tracking_settings_router, ) @@ -427,6 +430,9 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( router as pass_through_router, ) from litellm.proxy.policy_engine.policy_endpoints import router as policy_crud_router +from litellm.proxy.policy_engine.policy_resolve_endpoints import ( + router as policy_resolve_router, +) from litellm.proxy.prompts.prompt_endpoints import router as prompts_router from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router @@ -4603,11 +4609,16 @@ async def initialize( # noqa: PLR0915 elif litellm_log_setting.upper() == "DEBUG": import logging - from litellm._logging import verbose_proxy_logger, verbose_router_logger + from litellm._logging import ( + verbose_logger, + verbose_proxy_logger, + verbose_router_logger, + ) + verbose_logger.setLevel(level=logging.DEBUG) # set package log to debug verbose_router_logger.setLevel( level=logging.DEBUG - ) # set router logs to info + ) # set router logs to debug verbose_proxy_logger.setLevel( level=logging.DEBUG ) # set proxy logs to debug @@ -10372,7 +10383,17 @@ async def new_invitation( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Allow proxy admins and org/team admins (admin status from DB via get_user_object) + has_access = ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + or await _user_has_admin_privileges( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + if not has_access: raise HTTPException( status_code=400, detail={ @@ -10383,6 +10404,23 @@ async def new_invitation( }, ) + # Org/team admins can only invite users within their org/team + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + can_invite = await admin_can_invite_user( + target_user_id=data.user_id, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if not can_invite: + raise HTTPException( + status_code=400, + detail={ + "error": "You can only create invitations for users in your organization or team." + }, + ) + response = await create_invitation_for_user( data=data, user_api_key_dict=user_api_key_dict, @@ -10535,7 +10573,16 @@ async def invitation_delete( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Proxy admins can delete any invitation; org admins only their own + is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + is_other_admin = await _user_has_admin_privileges( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + if not is_proxy_admin and not is_other_admin: raise HTTPException( status_code=400, detail={ @@ -10546,6 +10593,24 @@ async def invitation_delete( }, ) + # Org admins can only delete invitations they created + if is_other_admin and not is_proxy_admin: + invitation = await prisma_client.db.litellm_invitationlink.find_unique( + where={"id": data.invitation_id} + ) + if invitation is None: + raise HTTPException( + status_code=400, + detail={"error": "Invitation id does not exist in the database."}, + ) + if invitation.created_by != user_api_key_dict.user_id: + raise HTTPException( + status_code=403, + detail={ + "error": "Organization admins can only delete invitations they created." + }, + ) + response = await prisma_client.db.litellm_invitationlink.delete( where={"id": data.invitation_id} ) @@ -11746,6 +11811,7 @@ app.include_router(analytics_router) app.include_router(guardrails_router) app.include_router(policy_router) app.include_router(policy_crud_router) +app.include_router(policy_resolve_router) app.include_router(search_tool_management_router) app.include_router(prompts_router) app.include_router(callback_management_endpoints_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 1750efed92c..37ed0182663 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -911,6 +911,7 @@ model LiteLLM_PolicyAttachmentTable { teams String[] @default([]) // Team aliases or patterns keys String[] @default([]) // Key aliases or patterns models String[] @default([]) // Model names or patterns + tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 6e49e4244e7..afbc57360e2 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1672,6 +1672,9 @@ async def ui_view_spend_logs( # noqa: PLR0915 model: Optional[str] = fastapi.Query( default=None, description="Filter logs by model" ), + model_id: Optional[str] = fastapi.Query( + default=None, description="Filter logs by model ID (litellm model deployment id)" + ), key_alias: Optional[str] = fastapi.Query( default=None, description="Filter logs by key alias" ), @@ -1684,6 +1687,14 @@ async def ui_view_spend_logs( # noqa: PLR0915 error_message: Optional[str] = fastapi.Query( default=None, description="Filter logs by error message (partial string match)" ), + sort_by: str = fastapi.Query( + default="startTime", + description="Sort logs by field: spend, total_tokens, startTime, or endTime", + ), + sort_order: Optional[str] = fastapi.Query( + default="desc", + description="Sort order: asc or desc", + ), ): """ View spend logs with pagination support. @@ -1715,6 +1726,23 @@ async def ui_view_spend_logs( # noqa: PLR0915 code=status.HTTP_400_BAD_REQUEST, ) + # Validate sort_by and sort_order + valid_sort_fields = {"spend", "total_tokens", "startTime", "endTime"} + if sort_by not in valid_sort_fields: + raise ProxyException( + message=f"Invalid sort_by: {sort_by}. Must be one of: {', '.join(sorted(valid_sort_fields))}", + type="bad_request", + param="sort_by", + code=status.HTTP_400_BAD_REQUEST, + ) + if sort_order is not None and sort_order.lower() not in {"asc", "desc"}: + raise ProxyException( + message=f"Invalid sort_order: {sort_order}. Must be one of: asc, desc", + type="bad_request", + param="sort_order", + code=status.HTTP_400_BAD_REQUEST, + ) + try: is_v2 = "/spend/logs/v2" in request.url.path formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"] @@ -1763,6 +1791,9 @@ async def ui_view_spend_logs( # noqa: PLR0915 if model is not None: where_conditions["model"] = model + if model_id is not None: + where_conditions["model_id"] = model_id + # Build metadata filters metadata_filters = [] if key_alias is not None: @@ -1824,6 +1855,11 @@ async def ui_view_spend_logs( # noqa: PLR0915 # Calculate skip value for pagination skip = (page - 1) * page_size + # Build order clause from sort_by and sort_order + order_column = sort_by + order_direction = (sort_order or "desc").lower() + order_clause = {order_column: order_direction} + # Get total count of records total_records = await prisma_client.db.litellm_spendlogs.count( where=where_conditions, @@ -1832,9 +1868,7 @@ async def ui_view_spend_logs( # noqa: PLR0915 # Get paginated data data = await prisma_client.db.litellm_spendlogs.find_many( where=where_conditions, - order={ - "startTime": "desc", - }, + order=order_clause, skip=skip, take=page_size, ) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index cb8b9ec0395..2f392c48e9d 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -641,7 +641,9 @@ def _sanitize_request_body_for_spend_logs_payload( return {k: _sanitize_value(v) for k, v in request_body.items()} -def _convert_to_json_serializable_dict(obj: Any) -> Any: +def _convert_to_json_serializable_dict( + obj: Any, visited: Optional[set] = None, max_depth: int = 20 +) -> Any: """ Convert object to JSON-serializable dict, handling Pydantic models safely. @@ -650,23 +652,55 @@ def _convert_to_json_serializable_dict(obj: Any) -> Any: Args: obj: Object to convert (dict, list, Pydantic model, or primitive) + visited: Set of object IDs to track circular references + max_depth: Maximum recursion depth to prevent infinite recursion Returns: JSON-serializable version of the object """ - if isinstance(obj, BaseModel): - # Use Pydantic's model_dump() instead of pickle - return obj.model_dump() - elif isinstance(obj, dict): - return {k: _convert_to_json_serializable_dict(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [_convert_to_json_serializable_dict(item) for item in obj] - elif hasattr(obj, "__dict__"): - # Handle objects with __dict__ attribute - return _convert_to_json_serializable_dict(obj.__dict__) - else: - # Primitives (str, int, float, bool, None) pass through - return obj + if max_depth <= 0: + # Return a placeholder if max depth is exceeded + return "" + + if visited is None: + visited = set() + + # Get the object's memory address to track visited objects + obj_id = id(obj) + if obj_id in visited: + # Circular reference detected, return placeholder + return "" + + # Only track mutable objects (dict, list, objects with __dict__) + if isinstance(obj, (dict, list)) or hasattr(obj, "__dict__"): + visited.add(obj_id) + + try: + if isinstance(obj, BaseModel): + # Use Pydantic's model_dump() instead of pickle + result = obj.model_dump() + # Recursively process the dumped dict + return _convert_to_json_serializable_dict(result, visited, max_depth - 1) + elif isinstance(obj, dict): + return { + k: _convert_to_json_serializable_dict(v, visited, max_depth - 1) + for k, v in obj.items() + } + elif isinstance(obj, list): + return [ + _convert_to_json_serializable_dict(item, visited, max_depth - 1) + for item in obj + ] + elif hasattr(obj, "__dict__"): + # Handle objects with __dict__ attribute + return _convert_to_json_serializable_dict(obj.__dict__, visited, max_depth - 1) + else: + # Primitives (str, int, float, bool, None) pass through + return obj + finally: + # Remove from visited set when done processing this object + if obj_id in visited: + visited.remove(obj_id) def _get_proxy_server_request_for_spend_logs_payload( diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index df298f7c448..2bff4e23c78 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -602,18 +602,53 @@ class LiteLLMCompletionResponsesConfig: } return None + @staticmethod + def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: + """ + Safely read a field from dict-like or attribute-based objects. + """ + if obj is None: + return default + + if isinstance(obj, dict): + return obj.get(key, default) + + getter = getattr(obj, "get", None) + if callable(getter): + try: + return getter(key, default) + except (TypeError, AttributeError): + pass + + return getattr(obj, key, default) + @staticmethod def _create_tool_call_chunk( tool_use_definition: Dict[str, Any], tool_call_id: str, index: int ) -> ChatCompletionToolCallChunk: """Create a ChatCompletionToolCallChunk from tool_use_definition.""" - function_raw = tool_use_definition.get("function") - function: Dict[str, Any] = function_raw if isinstance(function_raw, dict) else {} - tool_use_id_raw = tool_use_definition.get("id") + function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + tool_use_definition, "function" + ) + function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + function_raw, "name" + ) + function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + function_raw, "arguments" + ) + function: Dict[str, Any] = { + "name": function_name_raw or "", + "arguments": function_arguments_raw or "{}", + } + tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + tool_use_definition, "id" + ) tool_use_id: str = ( str(tool_use_id_raw) if tool_use_id_raw is not None else str(tool_call_id) ) - tool_use_type_raw = tool_use_definition.get("type") + tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + tool_use_definition, "type" + ) tool_use_type: str = ( str(tool_use_type_raw) if tool_use_type_raw is not None else "function" ) @@ -627,6 +662,63 @@ class LiteLLMCompletionResponsesConfig: index=index, ) + @staticmethod + def _normalize_tool_use_definition( + tool_use_definition: Any, tool_call_id: str + ) -> Optional[Dict[str, Any]]: + """ + Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk. + """ + if not tool_use_definition: + return None + + if isinstance(tool_use_definition, dict): + normalized_definition: Dict[str, Any] = dict(tool_use_definition) + else: + tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + tool_use_definition, "id" + ) + tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + tool_use_definition, "type" + ) + function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + tool_use_definition, "function" + ) + + # Object does not expose the expected tool_call fields. + if ( + tool_use_id_raw is None + and tool_use_type_raw is None + and function_raw is None + ): + return None + + normalized_definition = { + "id": tool_use_id_raw, + "type": tool_use_type_raw, + "function": function_raw, + } + + function_raw = normalized_definition.get("function") + if function_raw is not None and not isinstance(function_raw, dict): + function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + function_raw, "name" + ) + function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + function_raw, "arguments" + ) + if function_name_raw is not None or function_arguments_raw is not None: + normalized_definition["function"] = { + "name": function_name_raw, + "arguments": function_arguments_raw, + } + + normalized_definition["id"] = normalized_definition.get("id") or tool_call_id + normalized_definition["type"] = ( + normalized_definition.get("type") or "function" + ) + return normalized_definition + @staticmethod def _add_tool_call_to_assistant( assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk @@ -740,13 +832,19 @@ class LiteLLMCompletionResponsesConfig: tool_call_id, tools ) ) - - if _tool_use_definition: - if not isinstance(_tool_use_definition, dict): - _tool_use_definition = {} + + normalized_tool_use_definition = ( + LiteLLMCompletionResponsesConfig._normalize_tool_use_definition( + _tool_use_definition, tool_call_id + ) + ) + + if normalized_tool_use_definition: tool_call_chunk = ( LiteLLMCompletionResponsesConfig._create_tool_call_chunk( - _tool_use_definition, tool_call_id, len(tool_calls) + normalized_tool_use_definition, + tool_call_id, + len(tool_calls), ) ) LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant( @@ -1794,6 +1892,12 @@ class LiteLLMCompletionResponsesConfig: ): output_details_dict["text_tokens"] = completion_details.text_tokens + if ( + hasattr(completion_details, "image_tokens") + and completion_details.image_tokens is not None + ): + output_details_dict["image_tokens"] = completion_details.image_tokens + if output_details_dict: response_usage.output_tokens_details = OutputTokensDetails( **output_details_dict diff --git a/litellm/router.py b/litellm/router.py index 0f49cf90510..37fa3926b4d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -58,7 +58,6 @@ from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_metadata_variable_name_from_kwargs, ) -from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer @@ -619,11 +618,12 @@ class Router: self.retry_policy = RetryPolicy(**retry_policy) elif isinstance(retry_policy, RetryPolicy): self.retry_policy = retry_policy - verbose_router_logger.info( - "\033[32mRouter Custom Retry Policy Set:\n{}\033[0m".format( - self.retry_policy.model_dump(exclude_none=True) + if self.retry_policy is not None: + verbose_router_logger.info( + "\033[32mRouter Custom Retry Policy Set:\n{}\033[0m".format( + self.retry_policy.model_dump(exclude_none=True) + ) ) - ) self.model_group_retry_policy: Optional[ Dict[str, RetryPolicy] @@ -636,11 +636,12 @@ class Router: elif isinstance(allowed_fails_policy, AllowedFailsPolicy): self.allowed_fails_policy = allowed_fails_policy - verbose_router_logger.info( - "\033[32mRouter Custom Allowed Fails Policy Set:\n{}\033[0m".format( - self.allowed_fails_policy.model_dump(exclude_none=True) + if self.allowed_fails_policy is not None: + verbose_router_logger.info( + "\033[32mRouter Custom Allowed Fails Policy Set:\n{}\033[0m".format( + self.allowed_fails_policy.model_dump(exclude_none=True) + ) ) - ) self.alerting_config: Optional[AlertingConfig] = alerting_config @@ -1269,13 +1270,16 @@ class Router: if silent_model is not None: # Mirroring traffic to a secondary model - # Use shared thread pool for background calls - executor.submit( - self._silent_experiment_completion, - silent_model, - messages, - **kwargs, + # Use threading.Thread (not ThreadPoolExecutor) - executor.submit() + # requires pickling args, which fails when kwargs contain unpicklable + # objects (e.g. _thread.RLock from OTEL spans, loggers) in deployment. + thread = threading.Thread( + target=self._silent_experiment_completion, + args=(silent_model, messages), + kwargs=kwargs, + daemon=True, ) + thread.start() self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) kwargs.pop("silent_model", None) # Ensure it's not in kwargs either @@ -1924,6 +1928,17 @@ class Router: "deployment_model_name": deployment_model_name, } ) + + ## DEPLOYMENT-LEVEL TAGS + deployment_tags = deployment.get("litellm_params", {}).get("tags") + if deployment_tags: + existing_tags = kwargs[metadata_variable_name].get("tags") or [] + merged_tags = list(existing_tags) + for tag in deployment_tags: + if tag not in merged_tags: + merged_tags.append(tag) + kwargs[metadata_variable_name]["tags"] = merged_tags + kwargs["model_info"] = model_info kwargs["timeout"] = self._get_timeout( @@ -5925,9 +5940,20 @@ class Router: deployment.litellm_params.custom_llm_provider + "/" + _model_name ) + # For the shared backend key, strip custom pricing fields so that + # one deployment's pricing overrides don't pollute another + # deployment sharing the same backend model name. + # Each deployment's full pricing is already stored under its + # unique model_id above. + _custom_pricing_fields = CustomPricingLiteLLMParams.model_fields.keys() + _shared_model_info = { + k: v + for k, v in _model_info.items() + if k not in _custom_pricing_fields + } litellm.register_model( model_cost={ - _model_name: _model_info, + _model_name: _shared_model_info, } ) diff --git a/litellm/types/integrations/cloudzero.py b/litellm/types/integrations/cloudzero.py index aeda76aa5f9..e79500e08db 100644 --- a/litellm/types/integrations/cloudzero.py +++ b/litellm/types/integrations/cloudzero.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Optional +from typing import Any, Dict class CBFRecord(Dict[str, Any]): @@ -9,19 +9,23 @@ class CBFRecord(Dict[str, Any]): (e.g., 'time/usage_start', 'cost/cost'), we use a Dict base class rather than TypedDict to accommodate the special characters in field names. - Expected CBF fields: + Expected CBF fields (per LIT-1907): - time/usage_start: ISO-formatted UTC datetime (Optional[str]) - cost/cost: Billed cost (float) - - resource/id: CloudZero Resource Name (CZRN) (str) + - resource/id: Model name (str) - usage/amount: Numeric value of tokens consumed (int) - usage/units: Description of units, e.g., 'tokens' (str) - - resource/service: Maps to CZRN service-type, e.g., 'litellm' (str) - - resource/account: Maps to CZRN owner-account-id (entity_id) (str) + - resource/service: Model group (str) + - resource/account: api_key_alias|api_key_prefix (str) - resource/region: Maps to CZRN region, e.g., 'cross-region' (str) - - resource/usage_family: Maps to CZRN resource-type, e.g., 'llm-usage' (str) + - resource/usage_family: Provider (str) + - action/operation: Team ID (str) - lineitem/type: Standard usage line item, e.g., 'Usage' (str) - resource/tag:provider: CZRN provider component (str) - resource/tag:model: CZRN cloud-local-id component (model) (str) + - resource/tag:organization_alias: Organization alias if available (Optional[str]) + - resource/tag:project_alias: Project alias if available (Optional[str]) + - resource/tag:user_alias: User alias if available (Optional[str]) - resource/tag:{key}: Various resource tags for dimensions and metrics (Optional[str]) """ pass diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py index dc167667bc0..4a4cdaa2bae 100644 --- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -8,3 +8,4 @@ class UiDiscoveryEndpoints(BaseModel): proxy_base_url: Optional[str] auto_redirect_to_sso: bool admin_ui_disabled: bool + sso_configured: bool diff --git a/litellm/types/proxy/policy_engine/__init__.py b/litellm/types/proxy/policy_engine/__init__.py index bc54c3eb36b..42490c2eddc 100644 --- a/litellm/types/proxy/policy_engine/__init__.py +++ b/litellm/types/proxy/policy_engine/__init__.py @@ -19,6 +19,7 @@ from litellm.types.proxy.policy_engine.policy_types import ( PolicyScope, ) from litellm.types.proxy.policy_engine.resolver_types import ( + AttachmentImpactResponse, PolicyAttachmentCreateRequest, PolicyAttachmentDBResponse, PolicyAttachmentListResponse, @@ -30,6 +31,9 @@ from litellm.types.proxy.policy_engine.resolver_types import ( PolicyListDBResponse, PolicyListResponse, PolicyMatchContext, + PolicyMatchDetail, + PolicyResolveRequest, + PolicyResolveResponse, PolicyScopeResponse, PolicySummaryItem, PolicyTestResponse, @@ -75,4 +79,9 @@ __all__ = [ "PolicyAttachmentCreateRequest", "PolicyAttachmentDBResponse", "PolicyAttachmentListResponse", + # Resolve types + "PolicyResolveRequest", + "PolicyResolveResponse", + "PolicyMatchDetail", + "AttachmentImpactResponse", ] diff --git a/litellm/types/proxy/policy_engine/policy_types.py b/litellm/types/proxy/policy_engine/policy_types.py index 1c01f89e8b4..f221ba7e038 100644 --- a/litellm/types/proxy/policy_engine/policy_types.py +++ b/litellm/types/proxy/policy_engine/policy_types.py @@ -73,13 +73,15 @@ class PolicyScope(BaseModel): Used internally by PolicyAttachment to define WHERE a policy applies. Scope Fields: - | Field | What it matches | Wildcard support | - |--------|-----------------|----------------------| - | teams | Team aliases | *, healthcare-* | - | keys | Key aliases | *, dev-key-* | - | models | Model names | *, bedrock/*, gpt-* | + | Field | What it matches | Wildcard support | Default behavior | + |--------|-----------------|----------------------|---------------------| + | teams | Team aliases | *, healthcare-* | None → matches all | + | keys | Key aliases | *, dev-key-* | None → matches all | + | models | Model names | *, bedrock/*, gpt-* | None → matches all | + | tags | Key/team tags | *, health-*, prod-* | None → not checked | - If a field is None or empty, it defaults to matching everything (["*"]). + If teams/keys/models is None or empty, it defaults to matching everything (["*"]). + If tags is None or empty, the tag dimension is NOT checked (matches all). A request must match ALL specified scope fields for the attachment to apply. """ @@ -95,6 +97,10 @@ class PolicyScope(BaseModel): default=None, description="Model names or wildcard patterns. Use '*' for all models.", ) + tags: Optional[List[str]] = Field( + default=None, + description="Tag patterns to match against key/team tags. Supports wildcards (e.g., health-*).", + ) model_config = ConfigDict(extra="forbid") @@ -110,6 +116,14 @@ class PolicyScope(BaseModel): """Returns models list, defaulting to ['*'] if not specified.""" return self.models if self.models else ["*"] + def get_tags(self) -> List[str]: + """Returns tags list, defaulting to empty list if not specified. + + Unlike teams/keys/models, empty tags means 'do not check tags' + rather than 'match all'. This is because tags are opt-in scoping. + """ + return self.tags if self.tags else [] + # ───────────────────────────────────────────────────────────────────────────── # Policy Guardrails @@ -266,6 +280,10 @@ class PolicyAttachment(BaseModel): default=None, description="Model names or patterns this attachment applies to.", ) + tags: Optional[List[str]] = Field( + default=None, + description="Tag patterns this attachment applies to. Supports wildcards (e.g., health-*).", + ) model_config = ConfigDict(extra="forbid") @@ -281,6 +299,7 @@ class PolicyAttachment(BaseModel): teams=self.teams, keys=self.keys, models=self.models, + tags=self.tags, ) diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index 9488b8b0841..0c2c7336f8a 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -30,6 +30,10 @@ class PolicyMatchContext(BaseModel): default=None, description="Model name from the request.", ) + tags: Optional[List[str]] = Field( + default=None, + description="Tags from key/team metadata.", + ) model_config = ConfigDict(extra="forbid") @@ -65,6 +69,7 @@ class PolicyScopeResponse(BaseModel): teams: List[str] = Field(default_factory=list) keys: List[str] = Field(default_factory=list) models: List[str] = Field(default_factory=list) + tags: List[str] = Field(default_factory=list) class PolicyGuardrailsResponse(BaseModel): @@ -242,6 +247,10 @@ class PolicyAttachmentCreateRequest(BaseModel): default=None, description="Model names or patterns this attachment applies to.", ) + tags: Optional[List[str]] = Field( + default=None, + description="Tag patterns this attachment applies to. Supports wildcards (e.g., health-*).", + ) class PolicyAttachmentDBResponse(BaseModel): @@ -253,6 +262,7 @@ class PolicyAttachmentDBResponse(BaseModel): teams: List[str] = Field(default_factory=list, description="Team patterns.") keys: List[str] = Field(default_factory=list, description="Key patterns.") models: List[str] = Field(default_factory=list, description="Model patterns.") + tags: List[str] = Field(default_factory=list, description="Tag patterns.") created_at: Optional[datetime] = Field( default=None, description="When the attachment was created." ) @@ -274,3 +284,81 @@ class PolicyAttachmentListResponse(BaseModel): default_factory=list, description="List of policy attachments." ) total_count: int = Field(default=0, description="Total number of attachments.") + + +# ───────────────────────────────────────────────────────────────────────────── +# Policy Resolve Types +# ───────────────────────────────────────────────────────────────────────────── + + +class PolicyResolveRequest(BaseModel): + """Request body for resolving effective policies/guardrails for a context.""" + + team_alias: Optional[str] = Field( + default=None, description="Team alias to resolve for." + ) + key_alias: Optional[str] = Field( + default=None, description="Key alias to resolve for." + ) + model: Optional[str] = Field( + default=None, description="Model name to resolve for." + ) + tags: Optional[List[str]] = Field( + default=None, description="Tags to resolve for." + ) + + +class PolicyMatchDetail(BaseModel): + """Details about why a specific policy matched.""" + + policy_name: str = Field(description="Name of the matched policy.") + matched_via: str = Field( + description="How the policy was matched (e.g., 'tag:healthcare', 'team:health-team', 'scope:*')." + ) + guardrails_added: List[str] = Field( + default_factory=list, + description="Guardrails this policy contributes.", + ) + + +class PolicyResolveResponse(BaseModel): + """Response for resolving effective policies/guardrails for a context.""" + + effective_guardrails: List[str] = Field( + default_factory=list, + description="Final list of guardrails that would be applied.", + ) + matched_policies: List[PolicyMatchDetail] = Field( + default_factory=list, + description="Details about each matched policy and why it matched.", + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# Attachment Impact Estimation Types +# ───────────────────────────────────────────────────────────────────────────── + + +class AttachmentImpactResponse(BaseModel): + """Response for estimating the impact of a policy attachment.""" + + affected_keys_count: int = Field( + default=0, description="Number of keys that would be affected (named + unnamed)." + ) + affected_teams_count: int = Field( + default=0, description="Number of teams that would be affected (named + unnamed)." + ) + unnamed_keys_count: int = Field( + default=0, description="Number of affected keys without an alias." + ) + unnamed_teams_count: int = Field( + default=0, description="Number of affected teams without an alias." + ) + sample_keys: List[str] = Field( + default_factory=list, + description="Sample of affected key aliases (up to 10).", + ) + sample_teams: List[str] = Field( + default_factory=list, + description="Sample of affected team aliases (up to 10).", + ) diff --git a/litellm/utils.py b/litellm/utils.py index 6fdd2d88bca..0fa5436d98d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2506,6 +2506,16 @@ def _supports_factory(model: str, custom_llm_provider: Optional[str], key: str) if model_info.get(key, False) is True: return True elif model_info.get(key) is None: # don't check if 'False' explicitly set + # Fallback: when the provider-prefixed entry (e.g. + # "deepseek/deepseek-chat") exists but is missing a capability + # field, check the bare model-name entry (e.g. "deepseek-chat") + # which may carry the complete metadata. See #20885. + bare_model_key = _get_model_cost_key(model) + if bare_model_key is not None: + bare_entry = litellm.model_cost.get(bare_model_key) or {} + if bare_entry.get(key, False) is True: + return True + supported_by_provider = _supports_provider_info_factory( model, custom_llm_provider, key ) @@ -6140,6 +6150,13 @@ def validate_environment( # noqa: PLR0915 if ( "AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ + ) or ( + # IAM role, profile, or web identity auth don't require access keys + "AWS_ROLE_ARN" in os.environ + or "AWS_PROFILE" in os.environ + or "AWS_WEB_IDENTITY_TOKEN_FILE" in os.environ + or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" in os.environ # ECS task role + or "AWS_CONTAINER_CREDENTIALS_FULL_URI" in os.environ # ECS/Fargate full URI credential delivery ): keys_in_environment = True else: @@ -8243,6 +8260,8 @@ class ProviderConfigManager: return litellm.VolcEngineResponsesAPIConfig() elif litellm.LlmProviders.MANUS == provider: return litellm.ManusResponsesAPIConfig() + elif litellm.LlmProviders.PERPLEXITY == provider: + return litellm.PerplexityResponsesConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 815d29c7964..f6edcf7efd0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5848,6 +5848,19 @@ "output_cost_per_token": 7e-07, "supports_tool_choice": true }, + "azure_ai/kimi-k2.5": { + "input_cost_per_token": 6e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/kimi-k2-5-now-in-microsoft-foundry/4492321", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure_ai/ministral-3b": { "input_cost_per_token": 4e-08, "litellm_provider": "azure_ai", @@ -6091,6 +6104,39 @@ "output_cost_per_token": 2.4e-05, "supports_tool_choice": true }, + "bedrock/ap-northeast-1/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 7.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.03e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "bedrock/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 7.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.03e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "bedrock/moonshotai.kimi-k2.5": { + "input_cost_per_token": 7.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.03e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/ap-south-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 3.18e-06, "litellm_provider": "bedrock", @@ -6109,6 +6155,17 @@ "mode": "chat", "output_cost_per_token": 7.2e-07 }, + "bedrock/ap-south-1/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 7.1e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.94e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/ca-central-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 3.05e-06, "litellm_provider": "bedrock", @@ -6314,6 +6371,17 @@ "mode": "chat", "output_cost_per_token": 1.01e-06 }, + "bedrock/sa-east-1/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 7.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.03e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/us-east-1/1-month-commitment/anthropic.claude-instant-v1": { "input_cost_per_second": 0.011, "litellm_provider": "bedrock", @@ -6450,6 +6518,28 @@ "output_cost_per_token": 7e-07, "supports_tool_choice": true }, + "bedrock/us-east-1/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 6e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "bedrock/us-east-2/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 6e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/us-gov-east-1/amazon.nova-pro-v1:0": { "input_cost_per_token": 9.6e-07, "litellm_provider": "bedrock", @@ -6856,6 +6946,17 @@ "output_cost_per_token": 7e-07, "supports_tool_choice": true }, + "bedrock/us-west-2/moonshotai.kimi-k2-thinking": { + "input_cost_per_token": 6e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, "bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, "cache_read_input_token_cost": 8e-08, @@ -8945,6 +9046,43 @@ } ] }, + "dashscope/qwen3-max": { + "litellm_provider": "dashscope", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, "dashscope/qwq-plus": { "input_cost_per_token": 8e-07, "litellm_provider": "dashscope", @@ -10618,14 +10756,22 @@ "input_cost_per_token": 2.8e-07, "input_cost_per_token_cache_hit": 2.8e-08, "litellm_provider": "deepseek", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 4.2e-07, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, "supports_tool_choice": true }, "deepseek/deepseek-coder": { @@ -10662,16 +10808,24 @@ "input_cost_per_token": 2.8e-07, "input_cost_per_token_cache_hit": 2.8e-08, "litellm_provider": "deepseek", - "max_input_tokens": 128000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 4.2e-07, + "source": "https://api-docs.deepseek.com/quick_start/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_assistant_prefill": true, - "supports_function_calling": true, + "supports_function_calling": false, + "supports_native_streaming": true, + "supports_parallel_function_calling": false, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false }, "deepseek/deepseek-v3": { "cache_creation_input_token_cost": 0.0, @@ -25614,6 +25768,66 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "perplexity/preset/pro-search": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_preset": true + }, + "perplexity/openai/gpt-4o": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/openai/gpt-4o-mini": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/openai/gpt-5.2": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": true + }, + "perplexity/anthropic/claude-3-5-sonnet-20241022": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/anthropic/claude-3-5-haiku-20241022": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/google/gemini-2.0-flash-exp": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/google/gemini-2.0-flash-thinking-exp": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": true + }, + "perplexity/xai/grok-2-1212": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, + "perplexity/xai/grok-2-vision-1212": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false + }, "publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": { "input_cost_per_token": 0.0, "litellm_provider": "publicai", diff --git a/pyproject.toml b/pyproject.toml index f26e49093de..acb8bc2ada3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.81.9" +version = "1.81.10" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -61,7 +61,7 @@ boto3 = { version = "1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "0.4.33", optional = true} +litellm-proxy-extras = {version = "0.4.34", optional = true} rich = {version = "13.7.1", optional = true} litellm-enterprise = {version = "0.1.31", optional = true} diskcache = {version = "^5.6.1", optional = true} @@ -175,7 +175,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.81.9" +version = "1.81.10" version_files = [ "pyproject.toml:^version" ] diff --git a/requirements.txt b/requirements.txt index f680de120c5..bd313b105e5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -32,7 +32,7 @@ async_generator==1.10.0 # for async ollama calls langfuse==2.59.7 # for langfuse self-hosted logging prometheus_client==0.20.0 # for /metrics endpoint on proxy ddtrace==2.19.0 # for advanced DD tracing / profiling -orjson==3.11.2 # fast /embedding responses +orjson==3.11.7 # fast /embedding responses polars==1.31.0 # for data processing apscheduler==3.10.4 # for resetting budget in background fastapi-sso==0.19.0 # admin UI, SSO @@ -55,7 +55,7 @@ sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests cryptography==44.0.1 tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.33 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.34 # for proxy extras - e.g. prisma migrations llm-sandbox==0.3.31 # for skill execution in sandbox ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env diff --git a/schema.prisma b/schema.prisma index 9a87a491cf7..4329f939a7b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -913,6 +913,7 @@ model LiteLLM_PolicyAttachmentTable { teams String[] @default([]) // Team aliases or patterns keys String[] @default([]) // Key aliases or patterns models String[] @default([]) // Model names or patterns + tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt diff --git a/tests/batches_tests/test_batch_rate_limits.py b/tests/batches_tests/test_batch_rate_limits.py index 776aba438c2..13241e94d51 100644 --- a/tests/batches_tests/test_batch_rate_limits.py +++ b/tests/batches_tests/test_batch_rate_limits.py @@ -389,3 +389,233 @@ async def test_batch_rate_limit_multiple_requests(): print(f" Error: {exc_info.value.detail}") finally: os.unlink(file_path_2) + + +@pytest.mark.asyncio() +@pytest.mark.skipif( + os.environ.get("OPENAI_API_KEY") is None, + reason="OPENAI_API_KEY not set - skipping integration test" +) +async def test_batch_rate_limiter_with_managed_files(): + """ + Test for GEN-2166: Verify batch rate limiter can read user files when managed files are enabled. + + This test ensures that: + 1. The batch rate limiter passes user_api_key_dict to afile_content() + 2. The managed files hook can verify file ownership correctly + 3. Rate limiting is enforced (not silently bypassed) + 4. No 403 Permission Denied errors occur for files owned by the user + """ + import tempfile + from unittest.mock import AsyncMock, MagicMock, patch + + CUSTOM_LLM_PROVIDER = "openai" + + # Setup: Create internal usage cache and rate limiter + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + + # Setup: Get batch rate limiter + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None, "Batch rate limiter should be available" + + # Setup: Create user API key with TPM = 500, RPM = 10 + test_user_id = "test-user-abc123" + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key-managed-files", + user_id=test_user_id, + tpm_limit=500, + rpm_limit=10, + ) + + print(f"\n=== Testing Batch Rate Limiter with Managed Files ===") + print(f"User ID: {test_user_id}") + + # Create a batch file with ~200 tokens + import json as json_lib + message = "This is a test message for batch rate limiting with managed files. " * 5 + requests = [] + for i in range(1, 4): + request_obj = { + "custom_id": f"request-{i}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": message}] + } + } + requests.append(json_lib.dumps(request_obj)) + + batch_content = "\n".join(requests) + + with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f: + f.write(batch_content) + file_path = f.name + + try: + # Step 1: Upload file to OpenAI (simulating user upload) + print("\n1. Uploading batch input file...") + file_obj = await litellm.acreate_file( + file=open(file_path, "rb"), + purpose="batch", + custom_llm_provider=CUSTOM_LLM_PROVIDER, + ) + print(f" ✓ File uploaded: {file_obj.id}") + await asyncio.sleep(1) # Give API time to process + + # Step 2: Mock managed files hook to simulate file ownership check + # In a real scenario, the managed files hook would check if the user owns the file + # For this test, we'll verify that user_api_key_dict is passed correctly + print("\n2. Testing rate limiter file access with user context...") + + # Track if user_api_key_dict was passed to afile_content + original_afile_content = litellm.afile_content + user_context_passed = {"value": False} + + async def mock_afile_content(*args, **kwargs): + # Check if user_api_key_dict was passed + if "user_api_key_dict" in kwargs and kwargs["user_api_key_dict"] is not None: + user_context_passed["value"] = True + print(f" ✓ user_api_key_dict passed to afile_content") + print(f" User ID: {kwargs['user_api_key_dict'].user_id}") + else: + print(f" ✗ user_api_key_dict NOT passed to afile_content (BUG!)") + + # Call original function + return await original_afile_content(*args, **kwargs) + + # Patch afile_content to track the call + with patch('litellm.afile_content', side_effect=mock_afile_content): + data = { + "model": "gpt-3.5-turbo", + "input_file_id": file_obj.id, + "custom_llm_provider": CUSTOM_LLM_PROVIDER, + } + + # Step 3: Submit batch and verify rate limiting works + print("\n3. Submitting batch with rate limiting...") + result = await batch_limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=dual_cache, + data=data, + call_type="acreate_batch", + ) + + tokens_used = result.get('_batch_token_count', 0) + requests_count = result.get('_batch_request_count', 0) + print(f" ✓ Batch submitted successfully") + print(f" Tokens counted: {tokens_used}") + print(f" Requests counted: {requests_count}") + print(f" Rate limit usage: {tokens_used}/500 TPM, {requests_count}/10 RPM") + + # Step 4: Verify user context was passed + print("\n4. Verifying fix for GEN-2166...") + assert user_context_passed["value"], ( + "FAILED: user_api_key_dict was not passed to afile_content(). " + "This means the bug GEN-2166 is not fixed!" + ) + print(" ✓ Fix verified: user_api_key_dict is correctly passed") + + # Step 5: Verify rate limiting is actually enforced (not bypassed) + print("\n5. Verifying rate limiting is enforced...") + assert tokens_used > 0, "Token count should be greater than 0" + assert requests_count > 0, "Request count should be greater than 0" + print(" ✓ Rate limiting is active (not silently bypassed)") + + print("\n=== Test Passed: GEN-2166 Fix Verified ===") + print("✓ Batch rate limiter can access user files") + print("✓ User context is correctly passed") + print("✓ Rate limiting is enforced") + print("✓ No silent failures") + + except HTTPException as e: + if e.status_code == 403: + pytest.fail( + f"FAILED: Got 403 Permission Denied error. " + f"This indicates the bug GEN-2166 is not fixed. " + f"Error: {e.detail}" + ) + else: + raise + except Exception as e: + pytest.fail(f"Unexpected error: {str(e)}") + finally: + os.unlink(file_path) + + +@pytest.mark.asyncio() +async def test_batch_rate_limiter_without_user_context(): + """ + Test that verifies the bug scenario from GEN-2166. + + When user_api_key_dict is NOT passed to count_input_file_usage(), + the function should still work for non-managed files, but would fail + for managed files (which is the bug we fixed). + + This test documents the expected behavior with and without user context. + """ + import tempfile + + CUSTOM_LLM_PROVIDER = "openai" + + # Setup + BATCH_LIMITER = _PROXY_BatchRateLimiter( + internal_usage_cache=None, + parallel_request_limiter=None, + ) + + # Create a simple batch file + batch_content = """{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}""" + + with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f: + f.write(batch_content) + file_path = f.name + + try: + # Upload file + file_obj = await litellm.acreate_file( + file=open(file_path, "rb"), + purpose="batch", + custom_llm_provider=CUSTOM_LLM_PROVIDER, + ) + await asyncio.sleep(1) + + # Test 1: Without user context (old behavior - would fail with managed files) + print("\n=== Test 1: count_input_file_usage WITHOUT user context ===") + try: + usage_without_context = await BATCH_LIMITER.count_input_file_usage( + file_id=file_obj.id, + custom_llm_provider=CUSTOM_LLM_PROVIDER, + user_api_key_dict=None, # Explicitly passing None + ) + print(f"✓ Works for non-managed files (tokens: {usage_without_context.total_tokens})") + print(" Note: Would fail with 403 for managed files (GEN-2166 bug)") + except Exception as e: + print(f"✗ Failed: {str(e)}") + + # Test 2: With user context (new behavior - works with managed files) + print("\n=== Test 2: count_input_file_usage WITH user context ===") + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user-123", + ) + + usage_with_context = await BATCH_LIMITER.count_input_file_usage( + file_id=file_obj.id, + custom_llm_provider=CUSTOM_LLM_PROVIDER, + user_api_key_dict=user_api_key_dict, # Passing user context + ) + print(f"✓ Works with user context (tokens: {usage_with_context.total_tokens})") + print(" Note: This fixes GEN-2166 for managed files") + + # Verify both return the same results + assert usage_with_context.total_tokens == usage_without_context.total_tokens + assert usage_with_context.request_count == usage_without_context.request_count + print("\n✓ Both methods return identical results for non-managed files") + + finally: + os.unlink(file_path) diff --git a/tests/code_coverage_tests/check_guardrail_apply_decorator.py b/tests/code_coverage_tests/check_guardrail_apply_decorator.py new file mode 100644 index 00000000000..18a86277aa9 --- /dev/null +++ b/tests/code_coverage_tests/check_guardrail_apply_decorator.py @@ -0,0 +1,126 @@ +""" +Test that all guardrail hooks with async def apply_guardrail use @log_guardrail_information decorator. + +This ensures consistent logging and observability across all guardrail implementations. +""" + +import ast +from pathlib import Path +from typing import List, Tuple + + +def find_apply_guardrail_methods(file_path: Path) -> List[Tuple[str, int, bool]]: + """ + Find all apply_guardrail methods and check if they have the decorator. + + Returns: + List of tuples: (class_name, line_number, has_decorator) + """ + with open(file_path, "r") as f: + content = f.read() + + try: + tree = ast.parse(content) + except SyntaxError: + return [] + + results = [] + + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + class_name = node.name + + # Check if this class has apply_guardrail method + for item in node.body: + if ( + isinstance(item, ast.AsyncFunctionDef) + and item.name == "apply_guardrail" + ): + # Check if it has the log_guardrail_information decorator + has_decorator = False + for decorator in item.decorator_list: + if ( + isinstance(decorator, ast.Name) + and decorator.id == "log_guardrail_information" + ): + has_decorator = True + break + + results.append((class_name, item.lineno, has_decorator)) + + return results + + +def test_guardrail_apply_decorator(): + """Test that all guardrail hooks with apply_guardrail have the decorator.""" + # Path to the guardrail hooks directory + guardrail_hooks_dir = ( + Path(__file__).parent.parent.parent + / "litellm" + / "proxy" + / "guardrails" + / "guardrail_hooks" + ) + + # Find all Python files in the guardrail hooks directory + python_files = list(guardrail_hooks_dir.rglob("*.py")) + + # Track violations + violations = [] + + for python_file in python_files: + # Skip __init__.py files and test files + if python_file.name == "__init__.py" or python_file.name.startswith("test_"): + continue + + # Skip base files and primitives + if python_file.name in ["base.py", "primitives.py", "patterns.py"]: + continue + + # Skip bedrock_guardrails.py - it implements logging differently via + # add_standard_logging_guardrail_information_to_request_data calls + # in make_bedrock_api_request method instead of using the decorator + if python_file.name == "bedrock_guardrails.py": + continue + + results = find_apply_guardrail_methods(python_file) + + for class_name, line_num, has_decorator in results: + if not has_decorator: + relative_path = python_file.relative_to( + Path(__file__).parent.parent.parent + ) + violations.append((relative_path, class_name, line_num)) + + # Assert no violations found + if violations: + print( + f"\nFound {len(violations)} guardrail hook(s) without @log_guardrail_information decorator:" + ) + print( + "\nAll guardrail hooks must use @log_guardrail_information decorator on their apply_guardrail method." + ) + print( + "This ensures consistent logging and observability across all guardrails.\n" + ) + + for file_path, class_name, line_num in violations: + print(f" - {file_path}:{line_num} ({class_name}.apply_guardrail)") + + print("\nTo fix, add the decorator:") + print( + " from litellm.integrations.custom_guardrail import log_guardrail_information" + ) + print(" ") + print(" @log_guardrail_information") + print(" async def apply_guardrail(self, ...):") + print(" ...") + + raise AssertionError( + f"Found {len(violations)} guardrail hook(s) without @log_guardrail_information decorator" + ) + + +if __name__ == "__main__": + test_guardrail_apply_decorator() + print("✓ All guardrail hooks have @log_guardrail_information decorator") diff --git a/tests/code_coverage_tests/check_licenses.py b/tests/code_coverage_tests/check_licenses.py index dc7cc7c54ac..f49f6807a02 100644 --- a/tests/code_coverage_tests/check_licenses.py +++ b/tests/code_coverage_tests/check_licenses.py @@ -1,7 +1,8 @@ #!/usr/bin/env python3 import sys -import pkg_resources + import requests +from packaging.requirements import Requirement from pathlib import Path import json from typing import Dict, List, Optional, Set, Tuple @@ -213,7 +214,7 @@ class LicenseChecker: try: with open(requirements_file) as f: requirements = [ - pkg_resources.Requirement.parse(line) + Requirement(line.strip()) for line in f if line.strip() and not line.startswith("#") ] @@ -225,8 +226,10 @@ class LicenseChecker: for req in requirements: try: - version = next(iter(req.specs))[1] if req.specs else None - except Exception: + version = ( + next(iter(req.specifier)).version if req.specifier else None + ) + except StopIteration: version = None if not self.check_package(req.name, version): diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index dc46f83366a..e6e9d761ad5 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -111,6 +111,7 @@ ddtrace: >=2.19.0 # Unknown license orjson: >=3.10.12 # Unknown license apscheduler: >=3.10.4 # Unknown license fastapi-sso: >=0.16.0 # Unknown license +filelock: >=3.20.0 # Unlicense (public domain) - https://unlicense.org / https://github.com/tox-dev/filelock pyjwt: >=2.9.0 # Unknown license python-multipart: >=0.0.18 # Unknown license pillow: >=11.0.0 # Unknown license @@ -123,6 +124,7 @@ opentelemetry-exporter-otlp: >=1.25.0 # Unknown license sentry_sdk: >=2.21.0 # Unknown license cryptography: >=43.0.1 # Unknown license tzdata: >=2025.1 # Unknown license +urllib3: >=2.0.0 # MIT license - https://github.com/urllib3/urllib3 python-dotenv: >=1.0.0 # Unknown license tiktoken: >=0.8.0 # Unknown license click: >=8.1.7 # Unknown license diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 71e7798b09e..d6bf1941a08 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -42,6 +42,7 @@ IGNORE_FUNCTIONS = [ "_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation. "_basic_json_schema_validate", # max depth set. "extract_text_from_a2a_message", # max depth set (default 10) to prevent infinite recursion in A2A message parsing. + "_convert_to_json_serializable_dict", # max depth set (default 20) and circular reference protection to prevent infinite recursion. ] diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index 4bea2255b56..dff444168c2 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -78,7 +78,7 @@ async def test_bedrock_apply_guardrail_blocked(): }, ) - # HTTPException must propagate as-is (not wrapped in a generic Exception) + # Test the apply_guardrail method propagates HTTPException (AWS error) to the client with pytest.raises(HTTPException) as exc_info: await guardrail.apply_guardrail( inputs={"texts": ["This is blocked content"]}, @@ -276,7 +276,7 @@ async def test_bedrock_apply_guardrail_filters_request_messages_when_flag_enable }, ) - with pytest.raises(HTTPException) as exc_info: + with pytest.raises(HTTPException, match="policy") as exc_info: await guardrail.apply_guardrail( inputs={"texts": ["blocked"]}, request_data=request_data, @@ -286,9 +286,9 @@ async def test_bedrock_apply_guardrail_filters_request_messages_when_flag_enable assert mock_api.called _, kwargs = mock_api.call_args assert kwargs["messages"] == [request_messages[-1]] - # HTTPException must propagate as-is (not wrapped) + # HTTPException from guardrail is propagated so the client gets the AWS message assert exc_info.value.status_code == 400 - assert "Violated guardrail policy" in str(exc_info.value.detail) + assert "policy" in str(exc_info.value.detail) def test_bedrock_guardrail_filters_latest_user_message_when_enabled(): diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 946c5ad1729..58efa854e7c 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -192,7 +192,7 @@ async def test_async_post_call_success_hook_for_unified_finetuning_job(): "model_id": "gpt-3.5-turbo-0613", } proxy_managed_files = _PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() + DualCache(), prisma_client=AsyncMock() ) data = { "user_api_key_dict": {"parent_otel_span": MagicMock()}, @@ -373,6 +373,91 @@ async def test_output_file_id_for_batch_retrieve(): assert not cast(LiteLLMBatch, response).output_file_id.startswith("file-") +@pytest.mark.asyncio +async def test_error_file_id_for_failed_batch(): + """ + Test that the error_file_id is properly managed when a batch fails + """ + from typing import cast + + from openai.types.batch import BatchRequestCounts + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.openai import OpenAIFileObject + from litellm.types.utils import LiteLLMBatch + + batch = LiteLLMBatch( + id="bGl0ZWxsbV9wcm94eTttb2RlbF9pZDoxMjM0NTY3OTtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz", + completion_window="24h", + created_at=1714508499, + endpoint="/v1/chat/completions", + input_file_id="file-abc123", + object="batch", + status="failed", + cancelled_at=None, + cancelling_at=None, + completed_at=None, + error_file_id="error-abc123", + errors=None, + expired_at=None, + expires_at=1714536634, + failed_at=None, + finalizing_at=None, + in_progress_at=None, + metadata=None, + output_file_id=None, + request_counts=BatchRequestCounts(completed=0, failed=0, total=0), + usage=None, + ) + + batch._hidden_params = { + "litellm_call_id": "test-call-id", + "api_base": "https://api.openai.com", + "model_id": "test-model-id", + "model_name": "gpt-4o", + "response_cost": 0.0, + "additional_headers": {}, + "litellm_model_name": "gpt-4o", + "unified_batch_id": "litellm_proxy;model_id:test-model-id;llm_batch_id:batch_abc123", + } + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=AsyncMock() + ) + + # Create a proper OpenAIFileObject for the error file + error_file_object = OpenAIFileObject( + id="error-abc123", + object="file", + bytes=1234, + created_at=1714508500, + filename="error.jsonl", + purpose="batch_output", + status="processed", + ) + + # Mock the afile_retrieve to simulate retrieving error file metadata + with patch("litellm.afile_retrieve", new_callable=AsyncMock) as mock_retrieve: + mock_retrieve.return_value = error_file_object + + user_api_key_dict = UserAPIKeyAuth( + user_id="test-user-123", + parent_otel_span=MagicMock() + ) + + response = await proxy_managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=user_api_key_dict, + response=batch, + ) + + # Verify that error_file_id was transformed to a managed file ID + assert cast(LiteLLMBatch, response).error_file_id is not None + assert not cast(LiteLLMBatch, response).error_file_id.startswith("error-") + # Verify it's a base64 encoded managed file ID + assert _is_base64_encoded_unified_file_id(cast(LiteLLMBatch, response).error_file_id) + + @pytest.mark.asyncio async def test_async_post_call_success_hook_twice_assert_no_unique_violation(): import asyncio @@ -849,7 +934,7 @@ async def test_check_file_ids_access_with_unified_file_ids(): Test that check_file_ids_access validates user access to managed file IDs. """ from litellm.proxy._types import UserAPIKeyAuth - + # Create a unified file ID unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFkMDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My" regular_file_id = "file-abc123" diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 0567f60ecfc..3b4abeeb82f 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -5,7 +5,7 @@ import logging import os import sys import traceback -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch sys.path.insert( @@ -25,7 +25,7 @@ import pytest import litellm import json import tempfile -from base_image_generation_test import BaseImageGenTest +from base_image_generation_test import BaseImageGenTest, TestCustomLogger import logging from litellm._logging import verbose_logger @@ -182,6 +182,86 @@ class TestAimlImageGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "aiml/flux-pro/v1.1"} + @pytest.mark.asyncio(scope="module") + @pytest.mark.flaky(retries=0) + async def test_basic_image_generation(self): + """Test basic image generation""" + from unittest.mock import AsyncMock, patch + + mock_aiml_response = { + "created": 1703658209, + "data": [{"url": "https://example.com/generated_image.png"}], + } + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_aiml_response + mock_response.text = json.dumps(mock_aiml_response) + mock_response.headers = {} + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_async_post, patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + ) as mock_sync_post: + mock_async_post.return_value = mock_response + mock_sync_post.return_value = mock_response + + try: + litellm._turn_on_debug() + custom_logger = TestCustomLogger() + litellm.logging_callback_manager._reset_all_callbacks() + litellm.callbacks = [custom_logger] + base_image_generation_call_args = self.get_base_image_generation_call_args() + litellm.set_verbose = True + # Pass dummy api_key so validate_environment passes; HTTP is mocked + response = await litellm.aimage_generation( + **base_image_generation_call_args, + prompt="A image of a otter", + api_key="test-key-mocked-no-credits-needed", + ) + print("FAL AI RESPONSE: ", response) + + await asyncio.sleep(1) + + # assert response._hidden_params["response_cost"] is not None + # assert response._hidden_params["response_cost"] > 0 + # print("response_cost", response._hidden_params["response_cost"]) + + logged_standard_logging_payload = custom_logger.standard_logging_payload + print("logged_standard_logging_payload", logged_standard_logging_payload) + assert logged_standard_logging_payload is not None + assert logged_standard_logging_payload["response_cost"] is not None + assert logged_standard_logging_payload["response_cost"] > 0 + import openai + from openai.types.images_response import ImagesResponse + + # print openai version + print("openai version=", openai.__version__) + + response_dict = dict(response) + if "usage" in response_dict: + response_dict["usage"] = dict(response_dict["usage"]) + print("response usage=", response_dict.get("usage")) + + assert response.data is not None # type guard for iteration (base fails here if None) + for d in response.data: + assert isinstance(d, Image) + print("data in response.data", d) + assert d.b64_json is not None or d.url is not None + except litellm.RateLimitError as e: + pass + except litellm.ContentPolicyViolationError: + pass # Azure randomly raises these errors - skip when they occur + except litellm.InternalServerError: + pass + except Exception as e: + if "Your task failed as a result of our safety system." in str(e): + pass + else: + pytest.fail(f"An exception occurred - {str(e)}") + + class TestGoogleImageGen(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "gemini/imagen-4.0-generate-001"} diff --git a/tests/litellm/test_batch_completion_models_all_responses.py b/tests/litellm/test_batch_completion_models_all_responses.py new file mode 100644 index 00000000000..2e96ada03f2 --- /dev/null +++ b/tests/litellm/test_batch_completion_models_all_responses.py @@ -0,0 +1,118 @@ +import concurrent.futures + +import litellm +from litellm.batch_completion.main import batch_completion_models_all_responses + + +def test_batch_completion_models_all_responses_submits_before_waiting(monkeypatch): + """ + Regression test for issue #20704. + Ensures all model calls are submitted to the thread pool before waiting on results. + """ + models = ["model-a", "model-b", "model-c"] + called_models = [] + + class _AssertingFuture: + def __init__(self, result, executor, expected_submissions): + self._result = result + self._executor = executor + self._expected_submissions = expected_submissions + + def result(self): + if self._executor.submit_count != self._expected_submissions: + raise AssertionError("Not all model calls were submitted before waiting") + return self._result + + class _RecordingThreadPoolExecutor: + def __init__(self, max_workers, *args, **kwargs): + self.max_workers = max_workers + self.submit_count = 0 + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def submit(self, fn, *args, **kwargs): + self.submit_count += 1 + result = fn(*args, **kwargs) + return _AssertingFuture( + result=result, + executor=self, + expected_submissions=len(models), + ) + + def _mock_completion(*args, model, **kwargs): + called_models.append(model) + return {"model": model} + + monkeypatch.setattr(litellm, "completion", _mock_completion) + monkeypatch.setattr( + concurrent.futures, "ThreadPoolExecutor", _RecordingThreadPoolExecutor + ) + + responses = batch_completion_models_all_responses( + models=models, + messages=[{"role": "user", "content": "hello"}], + ) + + assert sorted(called_models) == sorted(models) + assert len(responses) == len(models) + assert sorted(response["model"] for response in responses) == sorted(models) + + +def test_batch_completion_models_all_responses_continues_on_model_error(monkeypatch): + models = ["model-a", "model-error", "model-b"] + + def _mock_completion(*args, model, **kwargs): + if model == "model-error": + raise RuntimeError("simulated model failure") + return {"model": model} + + monkeypatch.setattr(litellm, "completion", _mock_completion) + + responses = batch_completion_models_all_responses( + models=models, + messages=[{"role": "user", "content": "hello"}], + ) + + assert len(responses) == 2 + assert sorted(response["model"] for response in responses) == ["model-a", "model-b"] + + +def test_batch_completion_models_all_responses_returns_empty_for_empty_models(monkeypatch): + called = False + + def _mock_completion(*args, model, **kwargs): + nonlocal called + called = True + return {"model": model} + + monkeypatch.setattr(litellm, "completion", _mock_completion) + + responses = batch_completion_models_all_responses( + models=[], + messages=[{"role": "user", "content": "hello"}], + ) + + assert responses == [] + assert called is False + + +def test_batch_completion_models_all_responses_accepts_single_model_string(monkeypatch): + called_models = [] + + def _mock_completion(*args, model, **kwargs): + called_models.append(model) + return {"model": model} + + monkeypatch.setattr(litellm, "completion", _mock_completion) + + responses = batch_completion_models_all_responses( + models="model-a", + messages=[{"role": "user", "content": "hello"}], + ) + + assert called_models == ["model-a"] + assert responses == [{"model": "model-a"}] diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 3fd908f86d7..1da380b57a2 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -728,3 +728,18 @@ def test_azure_with_content_safety_error(): assert e.provider_specific_fields["innererror"]["code"] == "ResponsibleAIPolicyViolation" assert e.provider_specific_fields["innererror"]["content_filter_result"]["violence"]["filtered"] is True assert e.provider_specific_fields["innererror"]["content_filter_result"]["violence"]["severity"] == "high" + + +def test_azure_openai_with_prompt_cache_key(): + """ + E2E test for Azure OpenAI with prompt cache key param on /chat/completions API. + """ + litellm._turn_on_debug() + response = litellm.completion( + model="azure/gpt-4.1-mini", + api_key=os.getenv("AZURE_API_KEY"), + api_base=os.getenv("AZURE_API_BASE"), + api_version="2024-12-01-preview", + messages=[{"role": "user", "content": "What is the weather in San Francisco?"}], + prompt_cache_key="test_streaming_azure_openai", + ) \ No newline at end of file diff --git a/tests/llm_translation/test_model_cost_map_resilience.py b/tests/llm_translation/test_model_cost_map_resilience.py new file mode 100644 index 00000000000..61e375eabeb --- /dev/null +++ b/tests/llm_translation/test_model_cost_map_resilience.py @@ -0,0 +1,291 @@ +""" +Tests for model cost map resilience. + +Simulates: +- A bad (invalid JSON) model cost map upstream +- A bad (empty/missing) backup model cost map +- Verifies litellm.completion() still works even with a broken cost map +- Verifies litellm.get_model_info() raises the expected error for unmapped models +- Verifies the integrity validation helper catches corrupted maps +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) +) + +import litellm +from litellm.litellm_core_utils.get_model_cost_map import ( + GetModelCostMap, + get_model_cost_map, +) + + +class TestCheckIsValidDict: + """Unit tests for _check_is_valid_dict.""" + + def test_should_reject_non_dict(self): + """Non-dict should fail.""" + assert GetModelCostMap._check_is_valid_dict("not a dict") is False + + def test_should_reject_empty_dict(self): + """Empty dict should fail.""" + assert GetModelCostMap._check_is_valid_dict({}) is False + + def test_should_reject_list(self): + """List should fail.""" + assert GetModelCostMap._check_is_valid_dict([1, 2, 3]) is False + + def test_should_reject_none(self): + """None should fail.""" + assert GetModelCostMap._check_is_valid_dict(None) is False + + def test_should_accept_non_empty_dict(self): + """Non-empty dict should pass.""" + assert GetModelCostMap._check_is_valid_dict({"model": {}}) is True + + +class TestCheckModelCountNotReduced: + """Unit tests for _check_model_count_not_reduced.""" + + def test_should_reject_too_few_models(self): + """Fetched map with fewer models than min_model_count should fail.""" + small_map = {f"model-{i}": {} for i in range(5)} + assert ( + GetModelCostMap._check_model_count_not_reduced( + fetched_map=small_map, backup_model_count=0, min_model_count=10 + ) + is False + ) + + def test_should_reject_significant_shrinkage(self): + """Fetched map that shrunk >50% vs backup should fail.""" + fetched = {f"model-{i}": {} for i in range(40)} # 40% of 100 + assert ( + GetModelCostMap._check_model_count_not_reduced( + fetched_map=fetched, backup_model_count=100, min_model_count=10 + ) + is False + ) + + def test_should_accept_when_above_threshold(self): + """Fetched map at 60% of backup (above 50% threshold) should pass.""" + fetched = {f"model-{i}": {} for i in range(60)} + assert ( + GetModelCostMap._check_model_count_not_reduced( + fetched_map=fetched, backup_model_count=100, min_model_count=10 + ) + is True + ) + + def test_should_accept_growth(self): + """Fetched map larger than backup should pass.""" + fetched = {f"model-{i}": {} for i in range(120)} + assert ( + GetModelCostMap._check_model_count_not_reduced( + fetched_map=fetched, backup_model_count=100, min_model_count=10 + ) + is True + ) + + def test_should_accept_with_empty_backup(self): + """When backup is empty, only min_model_count matters.""" + fetched = {f"model-{i}": {} for i in range(15)} + assert ( + GetModelCostMap._check_model_count_not_reduced( + fetched_map=fetched, backup_model_count=0, min_model_count=10 + ) + is True + ) + + +class TestValidateModelCostMap: + """Unit tests for validate_model_cost_map (combines both checks).""" + + def test_should_reject_non_dict(self): + """Non-dict should fail at check 1.""" + assert GetModelCostMap.validate_model_cost_map(fetched_map="not a dict", backup_model_count=0) is False + + def test_should_reject_empty_map(self): + """Empty dict should fail at check 1.""" + assert GetModelCostMap.validate_model_cost_map(fetched_map={}, backup_model_count=0) is False + + def test_should_reject_significant_shrinkage(self): + """Should fail at check 2 (shrinkage).""" + fetched = {f"model-{i}": {} for i in range(40)} + assert ( + GetModelCostMap.validate_model_cost_map( + fetched_map=fetched, backup_model_count=100, min_model_count=10 + ) + is False + ) + + def test_should_accept_valid_map(self): + """Should pass both checks.""" + fetched = {f"model-{i}": {} for i in range(120)} + assert ( + GetModelCostMap.validate_model_cost_map( + fetched_map=fetched, backup_model_count=100, min_model_count=10 + ) + is True + ) + + def test_should_accept_equal_size_map(self): + """Equal size should pass both checks.""" + fetched = {f"model-{i}": {} for i in range(100)} + assert ( + GetModelCostMap.validate_model_cost_map( + fetched_map=fetched, backup_model_count=100, min_model_count=10 + ) + is True + ) + + +class TestGetModelCostMapFallback: + """Tests for get_model_cost_map fallback behavior with bad upstream.""" + + def test_should_fallback_to_backup_on_invalid_json(self): + """When upstream returns invalid JSON, should fall back to local backup.""" + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.side_effect = json.JSONDecodeError("bad json", "", 0) + + with patch("httpx.get", return_value=mock_response): + result = get_model_cost_map("https://fake-url.com/model_prices.json") + + # Should have fallen back to backup — backup always has models + assert isinstance(result, dict) + assert len(result) > 0 + + def test_should_fallback_to_backup_on_network_error(self): + """When upstream is unreachable, should fall back to local backup.""" + with patch("httpx.get", side_effect=Exception("Connection refused")): + result = get_model_cost_map("https://fake-url.com/model_prices.json") + + assert isinstance(result, dict) + assert len(result) > 0 + + def test_should_fallback_when_fetched_map_is_empty(self): + """When upstream returns valid JSON but empty dict, should fall back.""" + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = {} # empty map + + with patch("httpx.get", return_value=mock_response): + result = get_model_cost_map("https://fake-url.com/model_prices.json") + + # Should have fallen back to backup since empty map fails validation + assert isinstance(result, dict) + assert len(result) > 0 + + def test_should_fallback_when_fetched_map_shrinks_dramatically(self): + """When upstream returns far fewer models than backup, should fall back.""" + tiny_map = {f"model-{i}": {"litellm_provider": "test"} for i in range(11)} + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = tiny_map + + with patch("httpx.get", return_value=mock_response): + result = get_model_cost_map("https://fake-url.com/model_prices.json") + + # Backup has thousands of models; 11 is a massive shrinkage → fallback + assert len(result) > 11 + + def test_should_use_local_map_when_env_var_set(self): + """LITELLM_LOCAL_MODEL_COST_MAP=True should skip remote fetch entirely.""" + with patch.dict(os.environ, {"LITELLM_LOCAL_MODEL_COST_MAP": "True"}): + with patch("httpx.get") as mock_get: + result = get_model_cost_map( + "https://fake-url.com/model_prices.json" + ) + mock_get.assert_not_called() + + assert isinstance(result, dict) + assert len(result) > 0 + + +class TestBackupModelCostMapExists: + """Validates the local backup file is always present and valid.""" + + def test_should_have_backup_file(self): + """The backup model cost map must exist and be loadable.""" + backup = GetModelCostMap.load_local_model_cost_map() + assert isinstance(backup, dict) + assert len(backup) > 0, "Backup model cost map is empty" + + def test_should_have_minimum_models_in_backup(self): + """The backup must contain a reasonable number of models.""" + backup = GetModelCostMap.load_local_model_cost_map() + assert len(backup) > 100, ( + f"Backup has only {len(backup)} models, expected > 100" + ) + + +class TestBadHostedModelCostMap: + """ + Simulates the hosted model cost map being bad (invalid JSON / corrupted). + + When the hosted map is bad, get_model_cost_map() falls back to the local + backup. These tests verify that after fallback: + - get_model_info() still works for models in the backup + - litellm.completion() still works + """ + + def test_should_model_info_pass_after_bad_hosted_map(self): + """ + If the hosted map is bad, get_model_cost_map falls back to the local + backup. get_model_info should still work for models in the backup. + """ + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.side_effect = json.JSONDecodeError("bad json", "", 0) + + with patch("httpx.get", return_value=mock_response): + fallback_map = get_model_cost_map("https://fake-url.com/bad.json") + + original = litellm.model_cost + litellm.model_cost = fallback_map + try: + # gpt-4o is in every backup — should work fine + info = litellm.get_model_info("gpt-4o") + assert info is not None + assert info["input_cost_per_token"] > 0 + finally: + litellm.model_cost = original + + def test_should_completion_pass_after_bad_hosted_map(self): + """ + If the hosted map is bad, litellm.completion() should still work. + + Uses litellm's built-in mock_response param so the real completion + path is exercised (routing, cost calculator, logging) without + needing API credentials. + """ + # Simulate bad hosted map → fallback to backup + mock_http = MagicMock() + mock_http.raise_for_status = MagicMock() + mock_http.json.side_effect = json.JSONDecodeError("bad json", "", 0) + + with patch("httpx.get", return_value=mock_http): + fallback_map = get_model_cost_map("https://fake-url.com/bad.json") + + original = litellm.model_cost + litellm.model_cost = fallback_map + try: + # mock_response goes through the real completion path — + # routing, cost calculator, logging — but skips the HTTP call + response = litellm.completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "say hi"}], + mock_response="hello from mock", + ) + assert response is not None + assert response.choices[0].message.content == "hello from mock" + finally: + litellm.model_cost = original diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index e1c24ea46fe..0078483c734 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -3714,6 +3714,60 @@ def test_vertex_schema_test(): print(response) +def test_gemini_nullable_object_tool_schema_httpx(): + """ + Ensure nullable object tool params preserve nested properties in Vertex schema conversion. + """ + load_vertex_ai_credentials() + litellm._turn_on_debug() + + + tools = [{ + "type": "function", + "strict": True, + "function": { + "name": "create_support_ticket", + "description": "Create a paid user support ticket", + "parameters": { + "type": "object", + "additionalProperties": False, + "required": ["ticket_id", "customer_context"], + "properties": { + "ticket_id": { + "type": "string", + "description": "Unique identifier for the support ticket" + }, + "customer_context": { + "type": ["object", "null"], + "description": "Context about the paid customer, if available", + "additionalProperties": False, + "required": ["user_id", "plan"], + "properties": { + "user_id": { + "type": "string", + "description": "Internal user identifier" + }, + "plan": { + "type": "string", + "description": "Subscription plan name (e.g. pro, enterprise)" + } + } + } + } + } + } + }] + + response = litellm.completion( + model="vertex_ai/gemini-2.5-flash", + messages=[{"role": "user", "content": "call the tool"}], + tools=tools, + tool_choice="required", + ) + + print(response) + + def test_vertex_ai_response_id(): """Test that litellm preserves the response ID from Vertex AI's API for non-streaming responses""" from litellm.llms.custom_httpx.http_handler import HTTPHandler diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index b84fc22af01..d46a087eb73 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -311,7 +311,15 @@ def test_get_model_info_bedrock_models(): for commitment in potential_commitments: k = k.replace(f"{commitment}/", "") base_model = BedrockModelInfo.get_base_model(k) - base_model_info = litellm.model_cost[base_model] + # get_base_model() returns model id without "bedrock/" prefix; cost map keys use "bedrock/" + base_model_key = ( + base_model + if base_model in litellm.model_cost + else f"bedrock/{base_model}" + ) + if base_model_key not in litellm.model_cost: + continue + base_model_info = litellm.model_cost[base_model_key] for base_model_key, base_model_value in base_model_info.items(): if "invoke/" in k: continue diff --git a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py index b261daab6a7..e70d08b9008 100644 --- a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py +++ b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py @@ -452,16 +452,13 @@ async def test_redaction_with_metadata_completion_api(): litellm.callbacks = [test_custom_logger] # When metadata is passed, the system uses get_metadata_variable_name_from_kwargs - # to determine which field to check + # to determine which field to check. No headers means redaction should happen + # based on the global setting (litellm.turn_off_message_logging = True) response = await litellm.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "hi"}], mock_response="hello", - metadata={ - "headers": { - "litellm-disable-message-redaction": "true" - } - } + metadata={} ) await asyncio.sleep(1) diff --git a/tests/logging_callback_tests/test_otel_logging.py b/tests/logging_callback_tests/test_otel_logging.py index a0c78305e60..f0511e7d1ea 100644 --- a/tests/logging_callback_tests/test_otel_logging.py +++ b/tests/logging_callback_tests/test_otel_logging.py @@ -263,10 +263,18 @@ async def test_arize_phoenix_adds_openinference_kind_and_avoids_duplicate_litell Ensure Arize Phoenix spans include OpenInference span kind and do not create a duplicate litellm_request span when a proxy parent span is already active. """ + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor exporter.clear() litellm.logging_callback_manager._reset_all_callbacks() + # Set up a global TracerProvider so we can create valid spans + # This simulates the proxy server's TracerProvider + global_provider = TracerProvider() + global_provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace.set_tracer_provider(global_provider) + otel_logger = ArizePhoenixLogger(config=OpenTelemetryConfig(exporter=exporter)) litellm.callbacks = [otel_logger] litellm.success_callback = [] diff --git a/tests/mcp_tests/test_oauth2_mcp_config.yaml b/tests/mcp_tests/test_oauth2_mcp_config.yaml new file mode 100644 index 00000000000..c2704c5fe71 --- /dev/null +++ b/tests/mcp_tests/test_oauth2_mcp_config.yaml @@ -0,0 +1,14 @@ +model_list: + - model_name: fake-model + litellm_params: + model: openai/fake + api_key: fake-key + +mcp_servers: + test_oauth2_server: + url: "http://localhost:8765/mcp" + transport: "http" + auth_type: "oauth2" + client_id: "test-client" + client_secret: "test-secret" + token_url: "http://localhost:8765/oauth/token" diff --git a/tests/mcp_tests/test_openapi_spec_path_url.py b/tests/mcp_tests/test_openapi_spec_path_url.py new file mode 100644 index 00000000000..03e9db94967 --- /dev/null +++ b/tests/mcp_tests/test_openapi_spec_path_url.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from typing import Any, Dict + +import httpx +import pytest + +from litellm.proxy._experimental.mcp_server import openapi_to_mcp_generator as gen + + +class _FakeAsyncHTTPHandler: + """ + Minimal stand-in for the object returned by get_async_httpx_client(). + openapi_to_mcp_generator.load_openapi_spec_async() calls: + + client = get_async_httpx_client(...) + r = await client.get(url, timeout=30.0) + + So we must implement async get(). + """ + + def __init__(self, response: httpx.Response, expected_url: str): + self._response = response + self._expected_url = expected_url + self.calls = 0 + + async def get(self, request_url: str, timeout: float = 30.0): + self.calls += 1 + assert request_url == self._expected_url + assert timeout == 30.0 + return self._response + + +def test_load_openapi_spec_supports_http_url(monkeypatch: pytest.MonkeyPatch) -> None: + url = "http://example.local/openapi.json" + expected: Dict[str, Any] = { + "openapi": "3.0.0", + "info": {"title": "Test API", "version": "1.0.0"}, + "paths": {}, + } + + # httpx.Response must include a Request for raise_for_status() to work. + req = httpx.Request("GET", url) + resp = httpx.Response(status_code=200, json=expected, request=req) + + calls = {"get_async_httpx_client": 0} + handler_holder: Dict[str, Any] = {} + + def fake_get_async_httpx_client(*args, **kwargs): + calls["get_async_httpx_client"] += 1 + h = _FakeAsyncHTTPHandler(resp, expected_url=url) + handler_holder["handler"] = h + return h + + # Ensure shared/custom client path is used + monkeypatch.setattr(gen, "get_async_httpx_client", fake_get_async_httpx_client) + + # Fail loudly if someone reintroduces direct httpx.get() + def boom(*args, **kwargs): + raise AssertionError("Direct httpx.get() must not be used for URL spec loading") + + monkeypatch.setattr(httpx, "get", boom) + + spec = gen.load_openapi_spec(url) + + assert spec == expected + assert calls["get_async_httpx_client"] == 1 + assert handler_holder["handler"].calls == 1 + + +def test_load_openapi_spec_supports_local_file_path(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + expected: Dict[str, Any] = { + "openapi": "3.0.0", + "info": {"title": "Local API", "version": "1.0.0"}, + "paths": {}, + } + + p = tmp_path / "openapi.json" + p.write_text( + '{"openapi":"3.0.0","info":{"title":"Local API","version":"1.0.0"},"paths":{}}', + encoding="utf-8", + ) + + # For local files, shared client must NOT be used. + def boom_client(*args, **kwargs): + raise AssertionError("get_async_httpx_client() must not be called for local file paths") + + monkeypatch.setattr(gen, "get_async_httpx_client", boom_client) + + spec = gen.load_openapi_spec(str(p)) + assert spec == expected + diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py index 0a581fb512d..afc68dc9d42 100644 --- a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py +++ b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py @@ -525,7 +525,6 @@ async def test_anthropic_messages_with_extra_headers(): # Set up test parameters messages = [{"role": "user", "content": "Hello, can you tell me a short joke?"}] extra_headers = { - "anthropic-beta": "very-custom-beta-value", "anthropic-version": "custom-version-for-test", } @@ -581,87 +580,87 @@ async def test_anthropic_messages_with_extra_headers(): return response -@pytest.mark.asyncio -async def test_bedrock_messages_api_header_forwarding(): - """ - Test that headers from kwargs (set by proxy's add_headers_to_llm_call_by_model_group) - are correctly passed to validate_anthropic_messages_environment for Bedrock Invoke API. +# @pytest.mark.asyncio +# async def test_bedrock_messages_api_header_forwarding(): +# """ +# Test that headers from kwargs (set by proxy's add_headers_to_llm_call_by_model_group) +# are correctly passed to validate_anthropic_messages_environment for Bedrock Invoke API. - This verifies that forward_client_headers_to_llm_api works for Bedrock Invoke API (Messages API). +# This verifies that forward_client_headers_to_llm_api works for Bedrock Invoke API (Messages API). - Issue: When calling Anthropic models via the Messages API, LiteLLM makes a call to - Bedrock's Invoke API, and custom headers were not being forwarded, even though - they worked correctly for Chat Completions API with Bedrock's Converse API. - """ - from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.router import GenericLiteLLMParams +# Issue: When calling Anthropic models via the Messages API, LiteLLM makes a call to +# Bedrock's Invoke API, and custom headers were not being forwarded, even though +# they worked correctly for Chat Completions API with Bedrock's Converse API. +# """ +# from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +# from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +# from litellm.types.router import GenericLiteLLMParams - handler = BaseLLMHTTPHandler() +# handler = BaseLLMHTTPHandler() - # Headers that would be set by the proxy when forward_client_headers_to_llm_api is configured - custom_headers = { - "X-Custom-Header": "CustomValue", - "X-Request-ID": "req-123", - } +# # Headers that would be set by the proxy when forward_client_headers_to_llm_api is configured +# custom_headers = { +# "X-Custom-Header": "CustomValue", +# "X-Request-ID": "req-123", +# } - # Mock the provider config - mock_provider_config = MagicMock() +# # Mock the provider config +# mock_provider_config = MagicMock() - # We'll check what headers are passed to this method - mock_provider_config.validate_anthropic_messages_environment.return_value = ( - {"Authorization": "Bearer test"}, - "https://bedrock-runtime.us-east-1.amazonaws.com/invoke" - ) - mock_provider_config.transform_anthropic_messages_request.return_value = {"model": "test"} - mock_provider_config.get_complete_url.return_value = "https://test.com" - mock_provider_config.sign_request.return_value = ({}, None) - mock_provider_config.transform_anthropic_messages_response.return_value = {"id": "test"} +# # We'll check what headers are passed to this method +# mock_provider_config.validate_anthropic_messages_environment.return_value = ( +# {"Authorization": "Bearer test"}, +# "https://bedrock-runtime.us-east-1.amazonaws.com/invoke" +# ) +# mock_provider_config.transform_anthropic_messages_request.return_value = {"model": "test"} +# mock_provider_config.get_complete_url.return_value = "https://test.com" +# mock_provider_config.sign_request.return_value = ({}, None) +# mock_provider_config.transform_anthropic_messages_response.return_value = {"id": "test"} - # Mock HTTP client to prevent actual network calls - with unittest.mock.patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") as mock_get_client: - mock_http_client = AsyncMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = {"id": "test", "content": []} - mock_response.text = "{}" - mock_http_client.post.return_value = mock_response - mock_get_client.return_value = mock_http_client +# # Mock HTTP client to prevent actual network calls +# with unittest.mock.patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") as mock_get_client: +# mock_http_client = AsyncMock() +# mock_response = MagicMock() +# mock_response.status_code = 200 +# mock_response.json.return_value = {"id": "test", "content": []} +# mock_response.text = "{}" +# mock_http_client.post.return_value = mock_response +# mock_get_client.return_value = mock_http_client - # Mock logging object - mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) - mock_logging_obj.model_call_details = {} +# # Mock logging object +# mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) +# mock_logging_obj.model_call_details = {} - # Call the handler with headers in kwargs - try: - await handler.async_anthropic_messages_handler( - model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", - messages=[{"role": "user", "content": "Hello"}], - anthropic_messages_provider_config=mock_provider_config, - anthropic_messages_optional_request_params={"max_tokens": 100}, - custom_llm_provider="bedrock", - litellm_params=GenericLiteLLMParams( - api_key="test-key", - aws_region_name="us-east-1" - ), - logging_obj=mock_logging_obj, - api_key="test-key", - stream=False, - kwargs={"headers": custom_headers} # Headers set by proxy - ) - except Exception: - pass # Ignore errors, we're only checking if headers were passed +# # Call the handler with headers in kwargs +# try: +# await handler.async_anthropic_messages_handler( +# model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", +# messages=[{"role": "user", "content": "Hello"}], +# anthropic_messages_provider_config=mock_provider_config, +# anthropic_messages_optional_request_params={"max_tokens": 100}, +# custom_llm_provider="bedrock", +# litellm_params=GenericLiteLLMParams( +# api_key="test-key", +# aws_region_name="us-east-1" +# ), +# logging_obj=mock_logging_obj, +# api_key="test-key", +# stream=False, +# kwargs={"headers": custom_headers} # Headers set by proxy +# ) +# except Exception: +# pass # Ignore errors, we're only checking if headers were passed - # Verify that validate_anthropic_messages_environment was called - assert mock_provider_config.validate_anthropic_messages_environment.called +# # Verify that validate_anthropic_messages_environment was called +# assert mock_provider_config.validate_anthropic_messages_environment.called - # Get the headers that were passed - call_args = mock_provider_config.validate_anthropic_messages_environment.call_args - passed_headers = call_args[1]["headers"] +# # Get the headers that were passed +# call_args = mock_provider_config.validate_anthropic_messages_environment.call_args +# passed_headers = call_args[1]["headers"] - # The custom headers from kwargs should be in the passed headers - assert "X-Custom-Header" in passed_headers or "x-custom-header" in passed_headers - assert "X-Request-ID" in passed_headers or "x-request-id" in passed_headers +# # The custom headers from kwargs should be in the passed headers +# assert "X-Custom-Header" in passed_headers or "x-custom-header" in passed_headers +# assert "X-Request-ID" in passed_headers or "x-request-id" in passed_headers @pytest.mark.asyncio diff --git a/tests/pass_through_unit_tests/test_bedrock_tool_use_beta_header.py b/tests/pass_through_unit_tests/test_bedrock_tool_use_beta_header.py index 635ace016fe..e36b2ce9a5a 100644 --- a/tests/pass_through_unit_tests/test_bedrock_tool_use_beta_header.py +++ b/tests/pass_through_unit_tests/test_bedrock_tool_use_beta_header.py @@ -40,30 +40,30 @@ async def test_bedrock_sonnet_4_5_with_advanced_tool_use_beta_header(): print(f"✅ Test passed! Response: {response}") -@pytest.mark.asyncio -async def test_bedrock_claude_3_5_with_advanced_tool_use_beta_header_filtered(): - """ - Simple E2E test: Call Bedrock Claude 3.5 with advanced-tool-use beta header. +# @pytest.mark.asyncio +# async def test_bedrock_claude_3_5_with_advanced_tool_use_beta_header_filtered(): +# """ +# Simple E2E test: Call Bedrock Claude 3.5 with advanced-tool-use beta header. - This should work because the beta header is filtered out by LiteLLM before - sending the request to Bedrock Invoke API. - """ +# This should work because the beta header is filtered out by LiteLLM before +# sending the request to Bedrock Invoke API. +# """ - response = await litellm.anthropic.messages.acreate( - model="bedrock/invoke/us.anthropic.claude-3-5-sonnet-20240620-v1:0", - messages=[{"role": "user", "content": "What is 2+2?"}], - max_tokens=100, - provider_specific_header={ - "custom_llm_provider": "bedrock", - "extra_headers": { - "anthropic-beta": "advanced-tool-use-2025-11-20", - }, - }, - ) +# response = await litellm.anthropic.messages.acreate( +# model="bedrock/invoke/us.anthropic.claude-3-5-sonnet-20240620-v1:0", +# messages=[{"role": "user", "content": "What is 2+2?"}], +# max_tokens=100, +# provider_specific_header={ +# "custom_llm_provider": "bedrock", +# "extra_headers": { +# "anthropic-beta": "advanced-tool-use-2025-11-20", +# }, +# }, +# ) - # Verify response - assert response is not None - assert "content" in response - print(f"✅ Test passed! Claude 3.5 response (beta header filtered): {response}") +# # Verify response +# assert response is not None +# assert "content" in response +# print(f"✅ Test passed! Claude 3.5 response (beta header filtered): {response}") diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 54b9e31a6da..9aecfe9886b 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1845,38 +1845,38 @@ def test_provider_specific_header_multi_provider(): } -@pytest.mark.parametrize( - "custom_llm_provider, expected_result", - [ - ("anthropic", {"anthropic-beta": "test"}), - ("bedrock", {"anthropic-beta": "test"}), - ("vertex_ai", {"anthropic-beta": "test"}), - ], -) -def test_provider_specific_header_in_request(custom_llm_provider, expected_result): - from litellm.types.utils import ProviderSpecificHeader - from litellm.llms.custom_httpx.http_handler import HTTPHandler - from unittest.mock import patch +# @pytest.mark.parametrize( +# "custom_llm_provider, expected_result", +# [ +# ("anthropic", {"anthropic-beta": "test"}), +# ("bedrock", {"anthropic-beta": "test"}), +# ("vertex_ai", {"anthropic-beta": "test"}), +# ], +# ) +# def test_provider_specific_header_in_request(custom_llm_provider, expected_result): +# from litellm.types.utils import ProviderSpecificHeader +# from litellm.llms.custom_httpx.http_handler import HTTPHandler +# from unittest.mock import patch - litellm.set_verbose = True - client = HTTPHandler() - with patch.object(client, "post", return_value=MagicMock()) as mock_post: - try: - litellm.completion( - model="anthropic/claude-3-5-sonnet-v2@20241022", - messages=[{"role": "user", "content": "Hello world"}], - provider_specific_header=ProviderSpecificHeader( - custom_llm_provider="anthropic", - extra_headers={"anthropic-beta": "test"}, - ), - client=client, - ) - except Exception as e: - print(f"Error: {e}") +# litellm.set_verbose = True +# client = HTTPHandler() +# with patch.object(client, "post", return_value=MagicMock()) as mock_post: +# try: +# litellm.completion( +# model="anthropic/claude-3-5-sonnet-v2@20241022", +# messages=[{"role": "user", "content": "Hello world"}], +# provider_specific_header=ProviderSpecificHeader( +# custom_llm_provider="anthropic", +# extra_headers={"anthropic-beta": "test"}, +# ), +# client=client, +# ) +# except Exception as e: +# print(f"Error: {e}") - mock_post.assert_called_once() - print(mock_post.call_args.kwargs["headers"]) - assert "anthropic-beta" in mock_post.call_args.kwargs["headers"] +# mock_post.assert_called_once() +# print(mock_post.call_args.kwargs["headers"]) +# assert "anthropic-beta" in mock_post.call_args.kwargs["headers"] from litellm.proxy._types import LiteLLM_UserTable diff --git a/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py b/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py new file mode 100644 index 00000000000..a0dcf1c091c --- /dev/null +++ b/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py @@ -0,0 +1,169 @@ +""" +Tests that Arize Phoenix / Arize and the generic ``otel`` callback can +coexist, each sending spans to their own independent exporter. + +Covers the three root-cause fixes: +1. ArizePhoenixLogger / ArizeLogger create *dedicated* TracerProviders. +2. The ``otel`` dedup check does NOT match Arize subclasses. +3. Arize loggers do NOT overwrite ``proxy_server.open_telemetry_logger``. +""" + +import unittest +from unittest.mock import patch + +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_otel_logger(exporter: InMemorySpanExporter) -> OpenTelemetry: + """Create a generic ``otel`` callback backed by an in-memory exporter. + + We build a dedicated TracerProvider explicitly so the test is isolated + from whatever global provider state may exist. + """ + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + config = OpenTelemetryConfig(exporter=exporter) + return OpenTelemetry(config=config, callback_name="otel", tracer_provider=provider) + + +def _make_arize_phoenix_logger(exporter: InMemorySpanExporter): + """Create an ``arize_phoenix`` callback backed by an in-memory exporter. + + ArizePhoenixLogger._init_tracing creates its own TracerProvider, so we + pass the exporter via config and let it build the provider internally. + """ + from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger + + config = OpenTelemetryConfig(exporter=exporter) + return ArizePhoenixLogger(config=config, callback_name="arize_phoenix") + + +def _make_arize_logger(exporter: InMemorySpanExporter): + """Create an ``arize`` callback backed by an in-memory exporter. + + ArizeLogger._init_tracing creates its own TracerProvider, so we pass + the exporter via config and let it build the provider internally. + """ + from litellm.integrations.arize.arize import ArizeLogger + + config = OpenTelemetryConfig(exporter=exporter) + return ArizeLogger(config=config, callback_name="arize") + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +class TestIndependentTracerProviders(unittest.TestCase): + """Each integration must get its own TracerProvider so spans go to the right exporter.""" + + def test_otel_and_arize_phoenix_have_different_tracer_providers(self): + otel_exporter = InMemorySpanExporter() + phoenix_exporter = InMemorySpanExporter() + + otel_logger = _make_otel_logger(otel_exporter) + phoenix_logger = _make_arize_phoenix_logger(phoenix_exporter) + + # The tracers must come from different providers + assert otel_logger.tracer is not phoenix_logger.tracer + + def test_otel_and_arize_have_different_tracer_providers(self): + otel_exporter = InMemorySpanExporter() + arize_exporter = InMemorySpanExporter() + + otel_logger = _make_otel_logger(otel_exporter) + arize_logger = _make_arize_logger(arize_exporter) + + assert otel_logger.tracer is not arize_logger.tracer + + def test_arize_phoenix_and_arize_have_different_tracer_providers(self): + phoenix_exporter = InMemorySpanExporter() + arize_exporter = InMemorySpanExporter() + + phoenix_logger = _make_arize_phoenix_logger(phoenix_exporter) + arize_logger = _make_arize_logger(arize_exporter) + + assert phoenix_logger.tracer is not arize_logger.tracer + + +class TestSpansRoutedToCorrectExporter(unittest.TestCase): + """Spans created by each logger must land in its own exporter, not the other's.""" + + def test_spans_go_to_respective_exporters(self): + otel_exporter = InMemorySpanExporter() + phoenix_exporter = InMemorySpanExporter() + + otel_logger = _make_otel_logger(otel_exporter) + phoenix_logger = _make_arize_phoenix_logger(phoenix_exporter) + + # Create a span on each — SimpleSpanProcessor exports synchronously on end() + otel_span = otel_logger.tracer.start_span("otel_test_span") + otel_span.end() + + phoenix_span = phoenix_logger.tracer.start_span("phoenix_test_span") + phoenix_span.end() + + # Read spans *before* shutdown (shutdown clears the in-memory store) + otel_span_names = [s.name for s in otel_exporter.get_finished_spans()] + phoenix_span_names = [s.name for s in phoenix_exporter.get_finished_spans()] + + assert "otel_test_span" in otel_span_names + assert "phoenix_test_span" not in otel_span_names + + assert "phoenix_test_span" in phoenix_span_names + assert "otel_test_span" not in phoenix_span_names + + +class TestOtelDedupCheck(unittest.TestCase): + """The ``otel`` callback dedup must use exact type check, not isinstance.""" + + def test_arize_phoenix_logger_is_not_matched_by_otel_dedup(self): + from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger + + phoenix_logger = _make_arize_phoenix_logger(InMemorySpanExporter()) + + # isinstance would match — but type() must not + assert isinstance(phoenix_logger, OpenTelemetry) + assert type(phoenix_logger) is not OpenTelemetry + + def test_arize_logger_is_not_matched_by_otel_dedup(self): + from litellm.integrations.arize.arize import ArizeLogger + + arize_logger = _make_arize_logger(InMemorySpanExporter()) + + assert isinstance(arize_logger, OpenTelemetry) + assert type(arize_logger) is not OpenTelemetry + + def test_otel_logger_matches_own_dedup(self): + otel_logger = _make_otel_logger(InMemorySpanExporter()) + assert type(otel_logger) is OpenTelemetry + + +class TestProxyLoggerNotOverwritten(unittest.TestCase): + """Arize / Phoenix must not overwrite ``proxy_server.open_telemetry_logger``.""" + + @patch("litellm.proxy.proxy_server.open_telemetry_logger", None) + def test_arize_phoenix_does_not_set_proxy_otel_logger(self): + from litellm.proxy import proxy_server + + _make_arize_phoenix_logger(InMemorySpanExporter()) + assert proxy_server.open_telemetry_logger is None + + @patch("litellm.proxy.proxy_server.open_telemetry_logger", None) + def test_arize_does_not_set_proxy_otel_logger(self): + from litellm.proxy import proxy_server + + _make_arize_logger(InMemorySpanExporter()) + assert proxy_server.open_telemetry_logger is None + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_litellm/integrations/arize/test_arize_phoenix.py b/tests/test_litellm/integrations/arize/test_arize_phoenix.py index aa227fbff5e..129b35fb06a 100644 --- a/tests/test_litellm/integrations/arize/test_arize_phoenix.py +++ b/tests/test_litellm/integrations/arize/test_arize_phoenix.py @@ -1,5 +1,5 @@ import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -7,6 +7,7 @@ from litellm.integrations.arize.arize_phoenix import ( ArizePhoenixConfig, ArizePhoenixLogger, ) +from litellm.integrations.arize._utils import ArizeOTELAttributes class TestArizePhoenixConfig(unittest.TestCase): @@ -195,5 +196,63 @@ def test_get_arize_phoenix_config_expection_on_missing_api_key(monkeypatch, env_ +# --------------------------------------------------------------------------- +# Dynamic project naming from metadata +# --------------------------------------------------------------------------- + + +class TestGetDynamicProjectName: + """Tests for _get_dynamic_project_name extraction logic.""" + + def test_extracts_from_standard_logging_object_metadata(self): + kwargs = { + "standard_logging_object": { + "metadata": {"phoenix_project_name": "my-project"}, + } + } + assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) == "my-project" + + def test_extracts_from_litellm_params_metadata(self): + kwargs = { + "litellm_params": { + "metadata": {"phoenix_project_name": "sdk-project"}, + } + } + assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) == "sdk-project" + + def test_returns_none_when_no_metadata(self): + assert ArizePhoenixLogger._get_dynamic_project_name({}) is None + + def test_non_dict_standard_logging_object_does_not_raise(self): + """isinstance(dict) guard prevents AttributeError on non-dict payloads.""" + kwargs = {"standard_logging_object": "not-a-dict"} + assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) is None + + +class TestDynamicProjectNameOnSpan: + """set_arize_phoenix_attributes sets openinference.project.name on the span.""" + + @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-fallback"}, clear=False) + @patch("litellm.integrations.arize._utils.set_attributes") + def test_dynamic_name_sets_span_attribute(self, _mock_set_attrs): + span = MagicMock() + kwargs = { + "standard_logging_object": { + "metadata": {"phoenix_project_name": "dynamic-proj"}, + } + } + ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj=None) + + span.set_attribute.assert_called_once_with("openinference.project.name", "dynamic-proj") + + @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-project"}, clear=False) + @patch("litellm.integrations.arize._utils.set_attributes") + def test_falls_back_to_env_var_when_no_dynamic_name(self, _mock_set_attrs): + span = MagicMock() + ArizePhoenixLogger.set_arize_phoenix_attributes(span, {}, response_obj=None) + + span.set_attribute.assert_called_once_with("openinference.project.name", "env-project") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 95fa6ed8f60..3da9d9857c9 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -274,6 +274,26 @@ class TestOpenTelemetry(unittest.TestCase): self.assertEqual(config.deployment_environment, "production") self.assertEqual(config.model_id, "custom-service") + @patch.dict(os.environ, {}, clear=True) + def test_open_telemetry_config_auto_infer_otlp_http_when_endpoint_set(self): + """When endpoint is set but exporter is default 'console', auto-infer 'otlp_http'. + + This fixes an issue where UI-configured OTEL settings would default to console + output instead of sending traces to the configured endpoint. + See: https://github.com/BerriAI/litellm/issues/XXXX + """ + # When endpoint is specified without explicit exporter, should auto-infer otlp_http + config = OpenTelemetryConfig(endpoint="https://otel-collector.example.com:443") + self.assertEqual(config.exporter, "otlp_http") + + # When exporter is explicitly set to something other than console, should not override + config_grpc = OpenTelemetryConfig(exporter="grpc", endpoint="https://otel-collector.example.com:443") + self.assertEqual(config_grpc.exporter, "grpc") + + # When no endpoint is set, should keep console as default + config_no_endpoint = OpenTelemetryConfig() + self.assertEqual(config_no_endpoint.exporter, "console") + def wait_for_spans(self, exporter: InMemorySpanExporter, prefix: str): """Poll until we see at least one span with an attribute key starting with `prefix`.""" deadline = time.time() + self.POLL_TIMEOUT diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 980693aa73a..f566f91841d 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -12,6 +12,7 @@ sys.path.insert( from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_format_from_file_id, handle_any_messages_to_chat_completion_str_messages_conversion, + split_concatenated_json_objects, update_messages_with_model_file_ids, ) @@ -143,3 +144,58 @@ def test_convert_prefix_message_to_non_prefix_messages(): }, {"role": "assistant", "content": "value"}, ] + + +# ── split_concatenated_json_objects tests ── + + +def test_split_concatenated_json_single_object(): + """A single valid JSON object is returned as a one-element list.""" + result = split_concatenated_json_objects('{"location": "Boston"}') + assert result == [{"location": "Boston"}] + + +def test_split_concatenated_json_multiple_objects(): + """ + Multiple JSON objects concatenated without separators are split correctly. + This is the exact pattern from issue #20543 where Bedrock Claude Sonnet 4.5 + returns concatenated JSON in a single tool call arguments string. + """ + raw = ( + '{"command": ["curl", "-i", "http://localhost:9009"]}' + '{"command": ["curl", "-i", "http://localhost:9009/robots.txt"]}' + '{"command": ["curl", "-i", "http://localhost:9009/sitemap.xml"]}' + ) + result = split_concatenated_json_objects(raw) + assert len(result) == 3 + assert result[0] == {"command": ["curl", "-i", "http://localhost:9009"]} + assert result[1] == {"command": ["curl", "-i", "http://localhost:9009/robots.txt"]} + assert result[2] == {"command": ["curl", "-i", "http://localhost:9009/sitemap.xml"]} + + +def test_split_concatenated_json_with_whitespace(): + """Objects separated by whitespace are handled correctly.""" + raw = '{"a": 1} {"b": 2}\n{"c": 3}' + result = split_concatenated_json_objects(raw) + assert len(result) == 3 + assert result[0] == {"a": 1} + assert result[1] == {"b": 2} + assert result[2] == {"c": 3} + + +def test_split_concatenated_json_empty_string(): + """Empty or whitespace-only strings return an empty list.""" + assert split_concatenated_json_objects("") == [] + assert split_concatenated_json_objects(" ") == [] + + +def test_split_concatenated_json_non_dict_value(): + """Non-dict JSON values (e.g. arrays, strings) are replaced with {}.""" + result = split_concatenated_json_objects('[1, 2, 3]') + assert result == [{}] + + +def test_split_concatenated_json_invalid_raises(): + """Completely invalid JSON raises JSONDecodeError.""" + with pytest.raises(json.JSONDecodeError): + split_concatenated_json_objects("not json at all") diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index e87233a52a3..707b5bdc777 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( BAD_MESSAGE_ERROR_STR, BedrockConverseMessagesProcessor, BedrockImageProcessor, + _convert_to_bedrock_tool_call_invoke, ollama_pt, ) @@ -1590,3 +1591,153 @@ def test_bedrock_tools_unpack_defs_no_oom_with_nested_refs(): # Verify $defs have been removed (Bedrock doesn't support them) tool_schema = result[0]["toolSpec"].get("inputSchema", {}).get("json", {}) assert "$defs" not in tool_schema, "$defs should be removed after expansion" + + +# ── _convert_to_bedrock_tool_call_invoke tests ── + + +def test_bedrock_tool_call_invoke_normal_single_tool(): + """Normal single tool call with valid JSON arguments.""" + tool_calls = [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Boston, MA"}', + }, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 1 + assert result[0]["toolUse"]["toolUseId"] == "call_abc123" + assert result[0]["toolUse"]["name"] == "get_weather" + assert result[0]["toolUse"]["input"] == {"location": "Boston, MA"} + + +def test_bedrock_tool_call_invoke_empty_arguments(): + """Tool call with empty arguments produces an empty dict input.""" + tool_calls = [ + { + "id": "call_empty", + "type": "function", + "function": {"name": "do_something", "arguments": ""}, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 1 + assert result[0]["toolUse"]["input"] == {} + + +def test_bedrock_tool_call_invoke_concatenated_json(): + """ + Tool call whose arguments contain multiple concatenated JSON objects + (the bug from issue #20543) is split into separate Bedrock toolUse blocks. + + Bedrock Claude Sonnet 4.5 sometimes returns multiple tool call arguments + concatenated in a single string like: + '{"command":["curl",...]}{"command":["curl",...]}{"command":["curl",...]}' + """ + tool_calls = [ + { + "id": "tooluse_L7I3TewYAUhoheJZQEuwVN", + "type": "function", + "function": { + "name": "shell", + "arguments": ( + '{"command": ["curl", "-i", "http://localhost:9009", "-m", "10"]}' + '{"command": ["curl", "-i", "http://localhost:9009/robots.txt", "-m", "5"]}' + '{"command": ["curl", "-i", "http://localhost:9009/sitemap.xml", "-m", "5"]}' + ), + }, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + + # Should produce 3 separate toolUse blocks + assert len(result) == 3 + + # First block keeps original tool id + assert result[0]["toolUse"]["toolUseId"] == "tooluse_L7I3TewYAUhoheJZQEuwVN" + assert result[0]["toolUse"]["name"] == "shell" + assert result[0]["toolUse"]["input"] == { + "command": ["curl", "-i", "http://localhost:9009", "-m", "10"] + } + + # Subsequent blocks get suffixed ids + assert result[1]["toolUse"]["toolUseId"] == "tooluse_L7I3TewYAUhoheJZQEuwVN_1" + assert result[1]["toolUse"]["name"] == "shell" + assert result[1]["toolUse"]["input"] == { + "command": ["curl", "-i", "http://localhost:9009/robots.txt", "-m", "5"] + } + + assert result[2]["toolUse"]["toolUseId"] == "tooluse_L7I3TewYAUhoheJZQEuwVN_2" + assert result[2]["toolUse"]["name"] == "shell" + assert result[2]["toolUse"]["input"] == { + "command": ["curl", "-i", "http://localhost:9009/sitemap.xml", "-m", "5"] + } + + +def test_bedrock_tool_call_invoke_concatenated_json_with_cache_control(): + """ + When a tool call has cache_control AND concatenated JSON arguments, + the cachePoint block is appended after the last split block. + """ + tool_calls = [ + { + "id": "call_cached", + "type": "function", + "cache_control": {"type": "default"}, + "function": { + "name": "shell", + "arguments": '{"a": 1}{"b": 2}', + }, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + + # 2 toolUse blocks + 1 cachePoint block + assert len(result) == 3 + assert "toolUse" in result[0] + assert "toolUse" in result[1] + assert "cachePoint" in result[2] + + +def test_bedrock_tool_call_invoke_non_dict_arguments(): + """Arguments that parse to a non-dict (e.g. '""') produce empty dict input.""" + tool_calls = [ + { + "id": "call_non_dict", + "type": "function", + "function": {"name": "tool", "arguments": '""'}, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 1 + assert result[0]["toolUse"]["input"] == {} + + +def test_bedrock_tool_call_invoke_multiple_normal_tools(): + """Multiple separate tool calls (normal parallel calling) work correctly.""" + tool_calls = [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "NYC"}', + }, + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "LA"}', + }, + }, + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 2 + assert result[0]["toolUse"]["toolUseId"] == "call_1" + assert result[1]["toolUse"]["toolUseId"] == "call_2" diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py new file mode 100644 index 00000000000..d7df7823aee --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -0,0 +1,145 @@ +""" +Tests for litellm.litellm_core_utils.redact_messages.should_redact_message_logging + +Covers the proxy flow where headers arrive in litellm_params["metadata"]["headers"] +but litellm_params["litellm_metadata"] is None. +""" + +import pytest + +import litellm +from litellm.litellm_core_utils.redact_messages import should_redact_message_logging + + +@pytest.fixture(autouse=True) +def _reset_global_redaction(): + """Ensure the global setting is off for every test.""" + original = litellm.turn_off_message_logging + litellm.turn_off_message_logging = False + yield + litellm.turn_off_message_logging = original + + +def _make_model_call_details( + metadata_headers=None, + litellm_metadata=None, + metadata=None, + standard_callback_dynamic_params=None, +): + """Build a model_call_details dict that mimics real proxy/SDK flows.""" + litellm_params = {} + if metadata is not None: + litellm_params["metadata"] = metadata + elif metadata_headers is not None: + litellm_params["metadata"] = {"headers": metadata_headers} + else: + litellm_params["metadata"] = {} + + # get_litellm_params always sets this key (even when value is None) + litellm_params["litellm_metadata"] = litellm_metadata + + details = {"litellm_params": litellm_params} + if standard_callback_dynamic_params is not None: + details["standard_callback_dynamic_params"] = standard_callback_dynamic_params + return details + + +class TestShouldRedactMessageLogging: + """Unit tests for should_redact_message_logging().""" + + # ---- proxy flow: headers in metadata, litellm_metadata is None ---- + + def test_enable_redaction_via_x_header_proxy_flow(self): + """x-litellm-enable-message-redaction header should enable redaction + even when litellm_metadata is None (proxy path).""" + details = _make_model_call_details( + metadata_headers={"x-litellm-enable-message-redaction": "true"}, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is True + + def test_enable_redaction_via_old_header_proxy_flow(self): + """litellm-enable-message-redaction header should enable redaction + even when litellm_metadata is None (proxy path).""" + details = _make_model_call_details( + metadata_headers={"litellm-enable-message-redaction": "true"}, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is True + + def test_disable_redaction_via_header_proxy_flow(self): + """litellm-disable-message-redaction should suppress redaction + even when global setting is on, and litellm_metadata is None.""" + litellm.turn_off_message_logging = True + details = _make_model_call_details( + metadata_headers={"litellm-disable-message-redaction": "true"}, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is False + + # ---- SDK direct-call flow: headers in litellm_metadata ---- + + def test_enable_redaction_via_header_in_litellm_metadata(self): + """Headers inside litellm_metadata (SDK direct call) should work.""" + details = _make_model_call_details( + litellm_metadata={"headers": {"x-litellm-enable-message-redaction": "true"}}, + ) + assert should_redact_message_logging(details) is True + + # ---- no headers at all ---- + + def test_no_headers_defaults_to_global_off(self): + """Without headers, falls back to global setting (False).""" + details = _make_model_call_details( + metadata_headers=None, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is False + + def test_no_headers_global_on(self): + """Without headers, respects global turn_off_message_logging=True.""" + litellm.turn_off_message_logging = True + details = _make_model_call_details( + metadata_headers=None, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is True + + # ---- dynamic params take precedence ---- + + def test_dynamic_param_enables_redaction(self): + """Dynamic turn_off_message_logging=True should enable redaction.""" + details = _make_model_call_details( + metadata_headers={}, + litellm_metadata=None, + standard_callback_dynamic_params={"turn_off_message_logging": True}, + ) + assert should_redact_message_logging(details) is True + + def test_dynamic_param_false_overrides_header(self): + """Dynamic turn_off_message_logging=False should take precedence over enable header.""" + details = _make_model_call_details( + metadata_headers={"x-litellm-enable-message-redaction": "true"}, + litellm_metadata=None, + standard_callback_dynamic_params={"turn_off_message_logging": False}, + ) + assert should_redact_message_logging(details) is False + + # ---- non-dict metadata safety ---- + + def test_both_metadata_fields_none(self): + """When both litellm_metadata and metadata are None, should not raise.""" + details = _make_model_call_details( + metadata=None, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is False + + def test_both_metadata_fields_none_global_on(self): + """When both metadata fields are None but global is on, should still return True.""" + litellm.turn_off_message_logging = True + details = _make_model_call_details( + metadata=None, + litellm_metadata=None, + ) + assert should_redact_message_logging(details) is True diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index e3bd7d2bb31..57e0dd494e0 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -3,7 +3,6 @@ import os import sys import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../../../..") @@ -855,6 +854,92 @@ def test_anthropic_structured_output_beta_header(): ) +@pytest.mark.parametrize( + "model_name", + [ + "claude-opus-4-6-20250918", + "claude-opus-4.6-20250918", + "claude-opus-4-5-20251101", + "claude-opus-4.5-20251101", + ], +) +def test_opus_uses_native_structured_output(model_name): + """ + Test that Opus 4.5 and 4.6 models use native Anthropic structured outputs + (output_format) rather than the tool-based workaround. + """ + config = AnthropicConfig() + + response_format = { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"}, + }, + "required": ["name", "age"], + "additionalProperties": False, + }, + }, + } + + optional_params = config.map_openai_params( + non_default_params={"response_format": response_format}, + optional_params={}, + model=model_name, + drop_params=False, + ) + + # Should use output_format (native structured outputs) + assert "output_format" in optional_params + assert optional_params["output_format"]["type"] == "json_schema" + + # Should NOT create a tool-based workaround + assert "tools" not in optional_params + assert "tool_choice" not in optional_params + + # Should set json_mode + assert optional_params.get("json_mode") is True + + +def test_non_structured_output_model_uses_tool_workaround(): + """ + Test that models NOT in the native structured output list still use the + tool-based workaround for response_format. + """ + config = AnthropicConfig() + + response_format = { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + "required": ["result"], + "additionalProperties": False, + }, + }, + } + + optional_params = config.map_openai_params( + non_default_params={"response_format": response_format}, + optional_params={}, + model="claude-3-5-sonnet-20241022", + drop_params=False, + ) + + # Should NOT use output_format + assert "output_format" not in optional_params + + # Should use tool-based workaround + assert "tools" in optional_params + assert "tool_choice" in optional_params + + # ============ Tool Search Tests ============ diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 66d62aae1ec..80fd3ab698a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -185,7 +185,14 @@ def test_openai_model_with_thinking_converts_to_reasoning_effort(): # Verify reasoning_effort is set (converted from thinking) assert "reasoning_effort" in call_kwargs, "reasoning_effort should be passed to completion" - assert call_kwargs["reasoning_effort"] == "minimal", f"reasoning_effort should be 'minimal' for budget_tokens=1024, got {call_kwargs.get('reasoning_effort')}" + assert call_kwargs["reasoning_effort"] == { + "effort": "minimal", + "summary": "detailed", + }, f"reasoning_effort should request a reasoning summary for OpenAI responses API, got {call_kwargs.get('reasoning_effort')}" + + # Verify OpenAI thinking requests are routed to the Responses API + assert call_kwargs.get("model") == "responses/gpt-5.2" + # Verify thinking is NOT passed (non-Claude model) assert "thinking" not in call_kwargs, "thinking should NOT be passed for non-Claude models" diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index 8df35a37514..7be4d6dfcf2 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -30,6 +30,19 @@ class TestAzureOpenAIConfig: assert not config._is_response_format_supported_model("gpt-35-turbo") + def test_prompt_cache_key_supported(self): + """Test that 'prompt_cache_key' is in supported params for Azure OpenAI chat completion models. + + OpenAI's Chat Completions API supports prompt_cache_key for cache routing optimization. + """ + config = AzureOpenAIConfig() + supported_params = config.get_supported_openai_params("gpt-4.1-nano") + assert "prompt_cache_key" in supported_params + + supported_params = config.get_supported_openai_params("gpt-4.1") + assert "prompt_cache_key" in supported_params + + def test_map_openai_params_with_preview_api_version(): config = AzureOpenAIConfig() non_default_params = { diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index c91ef31bba5..d903d7c85f1 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -45,6 +45,58 @@ def test_azure_ai_validate_environment(): assert headers["Content-Type"] == "application/json" +def test_azure_ai_validate_environment_with_api_key(): + """ + Test that when api_key is provided, it is set in the api-key header + for Azure Foundry endpoints (.services.ai.azure.com). + """ + config = AzureAIStudioConfig() + headers = config.validate_environment( + headers={}, + model="Kimi-K2.5", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-api-key", + api_base="https://my-endpoint.services.ai.azure.com", + ) + assert headers["api-key"] == "test-api-key" + assert headers["Content-Type"] == "application/json" + + +def test_azure_ai_validate_environment_with_azure_ad_token(): + """ + Test that when no api_key is provided but Azure AD credentials are available, + the Authorization header is set with a Bearer token. + + Regression test for https://github.com/BerriAI/litellm/issues/20759 + """ + import litellm + + config = AzureAIStudioConfig() + with patch( + "litellm.llms.azure.common_utils.get_azure_ad_token", + return_value="fake-azure-ad-token", + ), patch( + "litellm.llms.azure.common_utils.get_secret_str", + return_value=None, + ), patch.object(litellm, "api_key", None), patch.object( + litellm, "azure_key", None + ): + headers = config.validate_environment( + headers={}, + model="Kimi-K2.5", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base="https://my-endpoint.services.ai.azure.com", + ) + assert headers.get("Authorization") == "Bearer fake-azure-ad-token" + assert "api-key" not in headers + assert headers["Content-Type"] == "application/json" + + def test_azure_ai_grok_stop_parameter_handling(): """ Test that Grok models properly handle stop parameter filtering in Azure AI Studio. diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index edfdeb08d82..d2fb45643de 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -281,153 +281,6 @@ def test_output_format_with_no_schema(): assert last_user_message["content"] == "Hello" -def test_advanced_tool_use_header_translation_for_opus_4_5(): - """ - Test that advanced-tool-use-2025-11-20 header is translated to Bedrock-specific headers - for Claude Opus 4.5. - - Regression test for: Claude Code sends advanced-tool-use header which needs to be - translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for Bedrock - Invoke API on Claude Opus 4.5. - - Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html - """ - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( - AmazonAnthropicClaudeMessagesConfig, - ) - - config = AmazonAnthropicClaudeMessagesConfig() - - messages = [ - {"role": "user", "content": "What's the weather like?"} - ] - - anthropic_messages_optional_request_params = { - "max_tokens": 100, - } - - # Simulate advanced-tool-use header from Claude Code - headers = { - "anthropic-beta": "advanced-tool-use-2025-11-20" - } - - # Test with Claude Opus 4.5 - result = config.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-5-20250514-v1:0", - messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - litellm_params={}, - headers=headers, - ) - - # Verify advanced-tool-use header was removed - assert "anthropic_beta" in result - beta_headers = result["anthropic_beta"] - assert "advanced-tool-use-2025-11-20" not in beta_headers, \ - "advanced-tool-use header should be removed for Bedrock" - - # Verify Bedrock-specific headers were added - assert "tool-search-tool-2025-10-19" in beta_headers, \ - "tool-search-tool-2025-10-19 should be added for Opus 4.5" - assert "tool-examples-2025-10-29" in beta_headers, \ - "tool-examples-2025-10-29 should be added for Opus 4.5" - - -def test_advanced_tool_use_header_filtered_for_non_opus_4_5(): - """ - Test that advanced-tool-use-2025-11-20 header is filtered out for models - that don't support tool search on Bedrock. - - Tool search is supported on: Claude Opus 4.5, Claude Sonnet 4.5 - Tool search is NOT supported on: Claude 3.5 Sonnet and earlier - """ - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( - AmazonAnthropicClaudeMessagesConfig, - ) - - config = AmazonAnthropicClaudeMessagesConfig() - - messages = [ - {"role": "user", "content": "What's the weather like?"} - ] - - anthropic_messages_optional_request_params = { - "max_tokens": 100, - } - - # Simulate advanced-tool-use header from Claude Code - headers = { - "anthropic-beta": "advanced-tool-use-2025-11-20" - } - - # Test with Claude 3.5 Sonnet (does NOT support tool search on Bedrock) - result = config.transform_anthropic_messages_request( - model="anthropic.claude-3-5-sonnet-20241022-v2:0", - messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - litellm_params={}, - headers=headers, - ) - - # Verify advanced-tool-use header was removed - beta_headers = result.get("anthropic_beta", []) - assert "advanced-tool-use-2025-11-20" not in beta_headers, \ - "advanced-tool-use header should be removed for Bedrock" - - # Verify Bedrock-specific headers were NOT added (only for Opus 4.5 and Sonnet 4.5) - assert "tool-search-tool-2025-10-19" not in beta_headers, \ - "tool-search-tool should not be added for models without tool search support" - assert "tool-examples-2025-10-29" not in beta_headers, \ - "tool-examples should not be added for models without tool search support" - - -def test_advanced_tool_use_header_translation_with_multiple_beta_headers(): - """ - Test that advanced-tool-use header translation works correctly when multiple - beta headers are present. - """ - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( - AmazonAnthropicClaudeMessagesConfig, - ) - - config = AmazonAnthropicClaudeMessagesConfig() - - messages = [ - {"role": "user", "content": "What's the weather like?"} - ] - - anthropic_messages_optional_request_params = { - "max_tokens": 100, - } - - # Multiple beta headers including advanced-tool-use - headers = { - "anthropic-beta": "claude-code-20250219,advanced-tool-use-2025-11-20,interleaved-thinking-2025-05-14" - } - - # Test with Claude Opus 4.5 - result = config.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-5-20250514-v1:0", - messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - litellm_params={}, - headers=headers, - ) - - beta_headers = result.get("anthropic_beta", []) - - # Verify advanced-tool-use was removed - assert "advanced-tool-use-2025-11-20" not in beta_headers - - # Verify Bedrock-specific headers were added - assert "tool-search-tool-2025-10-19" in beta_headers - assert "tool-examples-2025-10-29" in beta_headers - - # Verify other beta headers are preserved - assert "claude-code-20250219" in beta_headers - assert "interleaved-thinking-2025-05-14" in beta_headers - - def test_opus_4_5_model_detection(): """ Test that the _is_claude_opus_4_5 method correctly identifies Opus 4.5 models @@ -466,71 +319,71 @@ def test_opus_4_5_model_detection(): f"Should not detect {model} as Opus 4.5" -def test_structured_outputs_beta_header_filtered_for_bedrock_invoke(): - """ - Test that unsupported beta headers are filtered out for Bedrock Invoke API. +# def test_structured_outputs_beta_header_filtered_for_bedrock_invoke(): +# """ +# Test that unsupported beta headers are filtered out for Bedrock Invoke API. - Bedrock Invoke API only supports a specific whitelist of beta flags and returns - "invalid beta flag" error for others (e.g., structured-outputs, mcp-servers). - This test ensures unsupported headers are filtered while keeping supported ones. +# Bedrock Invoke API only supports a specific whitelist of beta flags and returns +# "invalid beta flag" error for others (e.g., structured-outputs, mcp-servers). +# This test ensures unsupported headers are filtered while keeping supported ones. - Fixes: https://github.com/BerriAI/litellm/issues/16726 - """ - config = AmazonAnthropicClaudeConfig() +# Fixes: https://github.com/BerriAI/litellm/issues/16726 +# """ +# config = AmazonAnthropicClaudeConfig() - messages = [{"role": "user", "content": "test"}] +# messages = [{"role": "user", "content": "test"}] - # Test 1: structured-outputs beta header (unsupported) - headers = {"anthropic-beta": "structured-outputs-2025-11-13"} +# # Test 1: structured-outputs beta header (unsupported) +# headers = {"anthropic-beta": "structured-outputs-2025-11-13"} - result = config.transform_request( - model="anthropic.claude-4-0-sonnet-20250514-v1:0", - messages=messages, - optional_params={}, - litellm_params={}, - headers=headers, - ) +# result = config.transform_request( +# model="anthropic.claude-4-0-sonnet-20250514-v1:0", +# messages=messages, +# optional_params={}, +# litellm_params={}, +# headers=headers, +# ) - # Verify structured-outputs beta is filtered out - anthropic_beta = result.get("anthropic_beta", []) - assert not any("structured-outputs" in beta for beta in anthropic_beta), \ - f"structured-outputs beta should be filtered, got: {anthropic_beta}" +# # Verify structured-outputs beta is filtered out +# anthropic_beta = result.get("anthropic_beta", []) +# assert not any("structured-outputs" in beta for beta in anthropic_beta), \ +# f"structured-outputs beta should be filtered, got: {anthropic_beta}" - # Test 2: mcp-servers beta header (unsupported - the main issue from #16726) - headers = {"anthropic-beta": "mcp-servers-2025-12-04"} +# # Test 2: mcp-servers beta header (unsupported - the main issue from #16726) +# headers = {"anthropic-beta": "mcp-servers-2025-12-04"} - result = config.transform_request( - model="anthropic.claude-4-0-sonnet-20250514-v1:0", - messages=messages, - optional_params={}, - litellm_params={}, - headers=headers, - ) +# result = config.transform_request( +# model="anthropic.claude-4-0-sonnet-20250514-v1:0", +# messages=messages, +# optional_params={}, +# litellm_params={}, +# headers=headers, +# ) - # Verify mcp-servers beta is filtered out - anthropic_beta = result.get("anthropic_beta", []) - assert not any("mcp-servers" in beta for beta in anthropic_beta), \ - f"mcp-servers beta should be filtered, got: {anthropic_beta}" +# # Verify mcp-servers beta is filtered out +# anthropic_beta = result.get("anthropic_beta", []) +# assert not any("mcp-servers" in beta for beta in anthropic_beta), \ +# f"mcp-servers beta should be filtered, got: {anthropic_beta}" - # Test 3: Mix of supported and unsupported beta headers - headers = {"anthropic-beta": "computer-use-2024-10-22,mcp-servers-2025-12-04,structured-outputs-2025-11-13"} +# # Test 3: Mix of supported and unsupported beta headers +# headers = {"anthropic-beta": "computer-use-2024-10-22,mcp-servers-2025-12-04,structured-outputs-2025-11-13"} - result = config.transform_request( - model="anthropic.claude-4-0-sonnet-20250514-v1:0", - messages=messages, - optional_params={}, - litellm_params={}, - headers=headers, - ) +# result = config.transform_request( +# model="anthropic.claude-4-0-sonnet-20250514-v1:0", +# messages=messages, +# optional_params={}, +# litellm_params={}, +# headers=headers, +# ) - # Verify only supported betas are kept - anthropic_beta = result.get("anthropic_beta", []) - assert not any("structured-outputs" in beta for beta in anthropic_beta), \ - f"structured-outputs beta should be filtered, got: {anthropic_beta}" - assert not any("mcp-servers" in beta for beta in anthropic_beta), \ - f"mcp-servers beta should be filtered, got: {anthropic_beta}" - assert any("computer-use" in beta for beta in anthropic_beta), \ - f"computer-use beta should be kept, got: {anthropic_beta}" +# # Verify only supported betas are kept +# anthropic_beta = result.get("anthropic_beta", []) +# assert not any("structured-outputs" in beta for beta in anthropic_beta), \ +# f"structured-outputs beta should be filtered, got: {anthropic_beta}" +# assert not any("mcp-servers" in beta for beta in anthropic_beta), \ +# f"mcp-servers beta should be filtered, got: {anthropic_beta}" +# assert any("computer-use" in beta for beta in anthropic_beta), \ +# f"computer-use beta should be kept, got: {anthropic_beta}" def test_output_format_removed_from_bedrock_invoke_request(): diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py index e7b6de29b6b..074a319a603 100644 --- a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py +++ b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py @@ -389,104 +389,4 @@ class TestAnthropicBetaHeaderSupport: assert "anthropic_beta" in additional_fields, ( "anthropic_beta SHOULD be added for Anthropic models with cross-region prefix." ) - assert "context-1m-2025-08-07" in additional_fields["anthropic_beta"] - - def test_messages_advanced_tool_use_translation_opus_4_5(self): - """Test that advanced-tool-use header is translated to Bedrock-specific headers for Opus 4.5. - - Regression test for: Claude Code sends advanced-tool-use-2025-11-20 header which needs - to be translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for - Bedrock Invoke API on Claude Opus 4.5. - - Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html - """ - config = AmazonAnthropicClaudeMessagesConfig() - headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} - - result = config.transform_anthropic_messages_request( - model="us.anthropic.claude-opus-4-5-20250514-v1:0", - messages=[{"role": "user", "content": "Test"}], - anthropic_messages_optional_request_params={"max_tokens": 100}, - litellm_params={}, - headers=headers - ) - - assert "anthropic_beta" in result - beta_headers = result["anthropic_beta"] - - # advanced-tool-use should be removed - assert "advanced-tool-use-2025-11-20" not in beta_headers, ( - "advanced-tool-use-2025-11-20 should be removed for Bedrock Invoke API" - ) - - # Bedrock-specific headers should be added for Opus 4.5 - assert "tool-search-tool-2025-10-19" in beta_headers, ( - "tool-search-tool-2025-10-19 should be added for Opus 4.5" - ) - assert "tool-examples-2025-10-29" in beta_headers, ( - "tool-examples-2025-10-29 should be added for Opus 4.5" - ) - - def test_messages_advanced_tool_use_translation_sonnet_4_5(self): - """Test that advanced-tool-use header is translated to Bedrock-specific headers for Sonnet 4.5. - - Regression test for: Claude Code sends advanced-tool-use-2025-11-20 header which needs - to be translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for - Bedrock Invoke API on Claude Sonnet 4.5. - - Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool - """ - config = AmazonAnthropicClaudeMessagesConfig() - headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} - - result = config.transform_anthropic_messages_request( - model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", - messages=[{"role": "user", "content": "Test"}], - anthropic_messages_optional_request_params={"max_tokens": 100}, - litellm_params={}, - headers=headers - ) - - assert "anthropic_beta" in result - beta_headers = result["anthropic_beta"] - - # advanced-tool-use should be removed - assert "advanced-tool-use-2025-11-20" not in beta_headers, ( - "advanced-tool-use-2025-11-20 should be removed for Bedrock Invoke API" - ) - - # Bedrock-specific headers should be added for Sonnet 4.5 - assert "tool-search-tool-2025-10-19" in beta_headers, ( - "tool-search-tool-2025-10-19 should be added for Sonnet 4.5" - ) - assert "tool-examples-2025-10-29" in beta_headers, ( - "tool-examples-2025-10-29 should be added for Sonnet 4.5" - ) - - def test_messages_advanced_tool_use_filtered_unsupported_model(self): - """Test that advanced-tool-use header is filtered out for models that don't support tool search. - - The translation to Bedrock-specific headers should only happen for models that - support tool search on Bedrock (Opus 4.5, Sonnet 4.5). - For other models, the advanced-tool-use header should just be removed. - """ - config = AmazonAnthropicClaudeMessagesConfig() - headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} - - # Test with Claude 3.5 Sonnet (does NOT support tool search on Bedrock) - result = config.transform_anthropic_messages_request( - model="us.anthropic.claude-3-5-sonnet-20241022-v2:0", - messages=[{"role": "user", "content": "Test"}], - anthropic_messages_optional_request_params={"max_tokens": 100}, - litellm_params={}, - headers=headers - ) - - beta_headers = result.get("anthropic_beta", []) - - # advanced-tool-use should be removed - assert "advanced-tool-use-2025-11-20" not in beta_headers - - # Bedrock-specific headers should NOT be added for unsupported models - assert "tool-search-tool-2025-10-19" not in beta_headers - assert "tool-examples-2025-10-29" not in beta_headers + assert "context-1m-2025-08-07" in additional_fields["anthropic_beta"] \ No newline at end of file diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 77eda432513..cf9fee6bacf 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -853,29 +853,99 @@ def test_role_assumption_ttl_calculation(): assert 3500 <= ttl <= 3600 # Allow some variance for test execution time -def test_role_assumption_error_handling(): +def test_role_assumption_access_denied_falls_back_when_same_role(): """ - Test that role assumption errors are properly propagated. + Test that when AssumeRole fails with AccessDenied AND the caller is confirmed + to already be running as the target role, we fall back to ambient credentials. """ base_aws_llm = BaseAWSLLM() - - # Mock the boto3 STS client to raise an exception + + # Mock the boto3 STS client to raise AccessDenied mock_sts_client = MagicMock() - mock_sts_client.assume_role.side_effect = Exception("AccessDenied: User is not authorized to perform sts:AssumeRole") - + mock_sts_client.assume_role.side_effect = Exception( + "An error occurred (AccessDenied) when calling the AssumeRole operation: " + "Roles may not be assumed by root accounts." + ) + + # Mock _auth_with_env_vars to return fallback credentials + mock_creds = MagicMock() + mock_creds.access_key = "fallback-access-key" + mock_creds.secret_key = "fallback-secret-key" + + with patch("boto3.client", return_value=mock_sts_client): + with patch.object( + base_aws_llm, "_auth_with_env_vars", return_value=(mock_creds, None) + ) as mock_env_auth: + # _is_already_running_as_role returns True => fallback allowed + with patch.object( + base_aws_llm, "_is_already_running_as_role", return_value=True + ): + credentials, ttl = base_aws_llm._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole", + aws_session_name="error-test-session", + ) + + # Should have fallen back to env vars + mock_env_auth.assert_called_once() + assert credentials.access_key == "fallback-access-key" + + +def test_role_assumption_access_denied_raises_when_different_role(): + """ + Test that when AssumeRole fails with AccessDenied but the caller is NOT + the same role, the error is re-raised (genuine permission failure). + """ + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.assume_role.side_effect = Exception( + "An error occurred (AccessDenied) when calling the AssumeRole operation: " + "User is not authorized to perform sts:AssumeRole" + ) + + with patch("boto3.client", return_value=mock_sts_client): + # _is_already_running_as_role returns False => do NOT fallback + with patch.object( + base_aws_llm, "_is_already_running_as_role", return_value=False + ): + with pytest.raises(Exception) as exc_info: + base_aws_llm._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole", + aws_session_name="error-test-session", + ) + + assert "AccessDenied" in str(exc_info.value) + + +def test_role_assumption_non_access_denied_error_propagated(): + """ + Test that non-AccessDenied errors from AssumeRole are still propagated. + """ + base_aws_llm = BaseAWSLLM() + + # Mock the boto3 STS client to raise a non-AccessDenied exception + mock_sts_client = MagicMock() + mock_sts_client.assume_role.side_effect = Exception( + "An error occurred (MalformedPolicyDocument) when calling the AssumeRole operation" + ) + with patch("boto3.client", return_value=mock_sts_client): - - # Should raise the exception with pytest.raises(Exception) as exc_info: base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, aws_session_token=None, - aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole", - aws_session_name="error-test-session" + aws_role_name="arn:aws:iam::1111111111111:role/BadPolicyRole", + aws_session_name="error-test-session", ) - - assert "AccessDenied" in str(exc_info.value) + + assert "MalformedPolicyDocument" in str(exc_info.value) def test_multiple_role_assumptions_in_sequence(): @@ -1195,3 +1265,251 @@ def test_converse_handler_external_id_extraction(): assert hasattr(mock_get_credentials, 'called_kwargs') assert "aws_external_id" in mock_get_credentials.called_kwargs assert mock_get_credentials.called_kwargs["aws_external_id"] == "TestExternalID123" + + +def test_is_already_running_as_role_irsa_same_role(): + """Test IRSA fast path: when AWS_ROLE_ARN matches target role.""" + base_aws_llm = BaseAWSLLM() + + with patch.dict(os.environ, { + "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole", + "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token", + }): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/MyRole" + ) is True + + +def test_is_already_running_as_role_irsa_different_role(): + """Test IRSA fast path: when AWS_ROLE_ARN does NOT match target role.""" + base_aws_llm = BaseAWSLLM() + + with patch.dict(os.environ, { + "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole", + "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token", + }): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::999999999999:role/OtherRole" + ) is False + + +def test_is_already_running_as_role_ecs_task_role(): + """Test ECS/EC2 path: GetCallerIdentity shows assumed-role matching target.""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + # Ensure no IRSA env vars + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/MyEcsTaskRole" + ) is True + + +def test_is_already_running_as_role_ecs_different_role(): + """Test ECS/EC2 path: GetCallerIdentity shows a different role.""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::999999999999:role/DifferentRole" + ) is False + + +def test_is_already_running_as_role_ecs_role_with_path(): + """Test ECS path with role that has a path prefix (e.g., /service-role/MyRole).""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + # Role ARN with path + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/service-role/MyEcsTaskRole" + ) is True + + +def test_is_already_running_as_role_get_caller_identity_fails(): + """Test that when GetCallerIdentity fails, we return False (don't crash).""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.side_effect = Exception("No credentials found") + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/SomeRole" + ) is False + + +def test_get_credentials_ecs_same_role_skips_assume_role(): + """ + End-to-end test: when running on ECS with the same role as aws_role_name, + get_credentials should use ambient credentials and NOT call AssumeRole. + """ + base_aws_llm = BaseAWSLLM() + + mock_creds = MagicMock() + mock_creds.access_key = "ecs-access-key" + mock_creds.secret_key = "ecs-secret-key" + mock_creds.token = "ecs-session-token" + + with patch.object( + base_aws_llm, + "_is_already_running_as_role", + return_value=True, + ): + with patch.object( + base_aws_llm, + "_auth_with_env_vars", + return_value=(mock_creds, None), + ) as mock_env_auth: + with patch.object( + base_aws_llm, + "_auth_with_aws_role", + ) as mock_role_auth: + credentials = base_aws_llm.get_credentials( + aws_role_name="arn:aws:iam::123456789012:role/MyEcsTaskRole", + aws_region_name="us-east-1", + ) + + # Should use env vars, NOT role assumption + mock_env_auth.assert_called_once() + mock_role_auth.assert_not_called() + assert credentials.access_key == "ecs-access-key" + + +def test_parse_arn_account_and_role_name(): + """Test the ARN parser helper for various ARN formats.""" + parse = BaseAWSLLM._parse_arn_account_and_role_name + + # Standard IAM role ARN + assert parse("arn:aws:iam::123456789012:role/MyRole") == ( + "aws", "123456789012", "MyRole" + ) + + # IAM role ARN with path + assert parse("arn:aws:iam::123456789012:role/service-role/MyRole") == ( + "aws", "123456789012", "MyRole" + ) + + # Assumed-role ARN (from GetCallerIdentity) + assert parse("arn:aws:sts::123456789012:assumed-role/MyRole/session-id") == ( + "aws", "123456789012", "MyRole" + ) + + # China partition + assert parse("arn:aws-cn:iam::123456789012:role/MyRole") == ( + "aws-cn", "123456789012", "MyRole" + ) + + # GovCloud partition + assert parse("arn:aws-us-gov:iam::123456789012:role/MyRole") == ( + "aws-us-gov", "123456789012", "MyRole" + ) + + # Invalid ARNs + assert parse("not-an-arn") is None + assert parse("arn:aws:iam::123456789012:user/MyUser") is None + assert parse("") is None + + +def test_is_already_running_as_role_cross_account_same_name(): + """ + Test that same role NAME in different accounts does NOT match. + This is the cross-account false-match prevention. + """ + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + # Caller is in account 111111111111 + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::111111111111:assumed-role/MyRole/session-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + # Target is same role name but in account 222222222222 + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::222222222222:role/MyRole" + ) is False + + +def test_is_already_running_as_role_cross_partition(): + """ + Test that same role name + account but different partition does NOT match. + """ + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyRole/session-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + # Same account and role but aws-cn partition + assert base_aws_llm._is_already_running_as_role( + "arn:aws-cn:iam::123456789012:role/MyRole" + ) is False + + +def test_is_already_running_as_role_invalid_target_arn(): + """ + Test that an unparseable target ARN returns False immediately. + """ + base_aws_llm = BaseAWSLLM() + + # Should return False without making any API calls + assert base_aws_llm._is_already_running_as_role("not-a-valid-arn") is False + + +def test_is_already_running_as_role_ssl_verify_passed(): + """ + Test that ssl_verify parameter is correctly passed to the STS client. + """ + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyRole/session-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client: + base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/MyRole", + ssl_verify="/path/to/ca-bundle.crt", + ) + mock_boto3_client.assert_called_once_with( + "sts", verify="/path/to/ca-bundle.crt" + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 65f08ef5021..c249bd9970c 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -140,6 +140,83 @@ async def test_ssl_verification_with_aiohttp_transport(): litellm.disable_aiohttp_transport = original_disable +@pytest.mark.asyncio +async def test_ssl_verification_with_shared_session(): + """ + Test that ssl_verify=False is respected even with shared sessions. + + This was a bug where shared sessions bypassed SSL configuration because + _create_aiohttp_transport returned immediately without passing ssl_verify + to the LiteLLMAiohttpTransport constructor. + + The fix stores ssl_verify in the transport and passes it per-request. + """ + import aiohttp + + # Ensure aiohttp transport is enabled for this test + original_disable = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = False + + try: + # Create a shared session (simulating what happens in production) + shared_session = aiohttp.ClientSession() + + try: + # Create transport with shared session and ssl_verify=False + transport = AsyncHTTPHandler._create_aiohttp_transport( + ssl_verify=False, + shared_session=shared_session, + ) + + # Verify the transport uses the shared session + assert transport.client is shared_session + + # Verify the SSL setting is stored in the transport for per-request use + assert transport._ssl_verify is False + finally: + await shared_session.close() + finally: + # Restore original setting + litellm.disable_aiohttp_transport = original_disable + + +@pytest.mark.asyncio +async def test_ssl_context_with_shared_session(): + """ + Test that ssl_context is respected even with shared sessions. + """ + import aiohttp + + # Ensure aiohttp transport is enabled for this test + original_disable = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = False + + try: + # Create a custom SSL context + custom_ssl_context = ssl.create_default_context() + + # Create a shared session + shared_session = aiohttp.ClientSession() + + try: + # Create transport with shared session and custom ssl_context + transport = AsyncHTTPHandler._create_aiohttp_transport( + ssl_context=custom_ssl_context, + shared_session=shared_session, + ) + + # Verify the transport uses the shared session + assert transport.client is shared_session + + # Verify the SSL context is stored in the transport for per-request use + assert transport._ssl_verify is custom_ssl_context + finally: + await shared_session.close() + finally: + # Restore original setting + litellm.disable_aiohttp_transport = original_disable + + @pytest.mark.asyncio async def test_aiohttp_transport_trust_env_setting(monkeypatch): """Test that trust_env setting is properly configured in aiohttp transport""" diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index f437b8405f7..f9b5b5fe29c 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -13,6 +13,7 @@ from unittest.mock import MagicMock, patch from litellm.llms.databricks.chat.transformation import ( DatabricksChatResponseIterator, DatabricksConfig, + _sanitize_empty_content, ) @@ -215,3 +216,45 @@ def test_chunk_parser_with_citation(): "end_char_index": 50, } } + + +def test_sanitize_empty_content_pops_none(): + message = {"role": "user", "content": None} + _sanitize_empty_content(message) + assert "content" not in message + + +def test_sanitize_empty_content_pops_empty_string(): + message = {"role": "user", "content": ""} + _sanitize_empty_content(message) + assert "content" not in message + + +def test_sanitize_empty_content_pops_single_empty_text_block(): + message = {"role": "user", "content": [{"type": "text", "text": ""}]} + _sanitize_empty_content(message) + assert "content" not in message + + +def test_sanitize_empty_content_filters_empty_blocks_keeps_non_empty(): + message = { + "role": "user", + "content": [ + {"type": "text", "text": ""}, + {"type": "text", "text": "Hello"}, + {"type": "text", "text": " "}, + ], + } + _sanitize_empty_content(message) + assert message["content"] == [{"type": "text", "text": "Hello"}] + + +def test_transform_messages_sanitizes_empty_content(): + config = DatabricksConfig() + messages = [ + {"role": "user", "content": [{"type": "text", "text": ""}]}, + {"role": "user", "content": "Hi"}, + ] + result = config._transform_messages(messages=messages, model="databricks-claude", is_async=False) + assert "content" not in result[0] + assert result[1]["content"] == "Hi" diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py index a9c4bead820..eed42519622 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -593,6 +593,113 @@ class TestOCICohereToolCalls: assert result.usage.total_tokens == 25 +class TestOCICoherePreambleOverride: + """Test Cohere system message handling via preambleOverride""" + + def test_single_system_message_sets_preamble_override(self): + """Test that a single system message is extracted into preambleOverride""" + config = OCIChatConfig() + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + ] + optional_params = {"oci_compartment_id": TEST_COMPARTMENT_ID} + + result = config.transform_request( + model="cohere.command-latest", + messages=messages, # type: ignore + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + chat_request = result["chatRequest"] + assert chat_request["preambleOverride"] == "You are a helpful assistant." + + def test_multiple_system_messages_combined(self): + """Test that multiple system messages are joined with newlines""" + config = OCIChatConfig() + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "system", "content": "Always respond in JSON."}, + {"role": "user", "content": "Hello"}, + ] + optional_params = {"oci_compartment_id": TEST_COMPARTMENT_ID} + + result = config.transform_request( + model="cohere.command-latest", + messages=messages, # type: ignore + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + chat_request = result["chatRequest"] + assert chat_request["preambleOverride"] == "You are a helpful assistant.\nAlways respond in JSON." + + def test_system_message_with_content_array(self): + """Test system message with list-style content (text blocks)""" + config = OCIChatConfig() + messages = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a coding assistant."}, + ], + }, + {"role": "user", "content": "Hello"}, + ] + optional_params = {"oci_compartment_id": TEST_COMPARTMENT_ID} + + result = config.transform_request( + model="cohere.command-latest", + messages=messages, # type: ignore + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + chat_request = result["chatRequest"] + assert chat_request["preambleOverride"] == "You are a coding assistant." + + def test_no_system_message_omits_preamble_override(self): + """Test that preambleOverride is omitted when there are no system messages""" + config = OCIChatConfig() + messages = [ + {"role": "user", "content": "Hello"}, + ] + optional_params = {"oci_compartment_id": TEST_COMPARTMENT_ID} + + result = config.transform_request( + model="cohere.command-latest", + messages=messages, # type: ignore + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + chat_request = result["chatRequest"] + assert "preambleOverride" not in chat_request + + def test_system_messages_excluded_from_chat_history(self): + """Test that system messages do not appear in chatHistory""" + config = OCIChatConfig() + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "First question"}, + {"role": "assistant", "content": "First answer"}, + {"role": "user", "content": "Second question"}, + ] + + chat_history = config.adapt_messages_to_cohere_standard(messages) + + # Should contain user and assistant only, no system + # Note: adapt_messages_to_cohere_standard excludes the last message + roles = [msg.role for msg in chat_history] + assert "SYSTEM" not in roles + assert roles == ["USER", "CHATBOT"] + + class TestOCICohereStreaming: """Test Cohere streaming functionality""" diff --git a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py index af6481a6cb0..02495106a84 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py @@ -10,7 +10,8 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) -from litellm.llms.ollama.chat.transformation import OllamaChatConfig +from litellm.llms.ollama.chat.transformation import OllamaChatConfig, OllamaChatCompletionResponseIterator + from litellm.types.llms.openai import AllMessageValues from litellm.utils import get_optional_params @@ -473,3 +474,130 @@ class TestOllamaToolCalling: # finish_reason should be "stop" (default behavior) assert result.choices[0].finish_reason == "stop" assert result.choices[0].message.tool_calls is None + + +class TestOllamaReasoningContentStreaming: + """Test that reasoning_content is properly extracted from all thinking chunks.""" + + def test_multiple_thinking_chunks_all_returned_as_reasoning_content(self): + """ + Test that more than 2 consecutive thinking chunks are all returned as reasoning_content. + + Previously, the code had a bug where finished_reasoning_content was set to True + after just 2 chunks with 'thinking', causing subsequent thinking content to be lost. + """ + iterator = OllamaChatCompletionResponseIterator( + streaming_response=iter([]), # Not used in chunk_parser + sync_stream=True, + ) + + # Simulate 5 consecutive chunks with 'thinking' content + thinking_chunks = [ + { + "model": "deepseek-r1", + "message": {"role": "assistant", "thinking": f"Thinking chunk {i}"}, + "done": False, + } + for i in range(1, 6) + ] + + # Process all thinking chunks + reasoning_contents = [] + for chunk in thinking_chunks: + result = iterator.chunk_parser(chunk) + rc = result.choices[0].delta.reasoning_content + reasoning_contents.append(rc) + + # ALL chunks should have reasoning_content, not just the first 2 + assert len(reasoning_contents) == 5 + assert reasoning_contents[0] == "Thinking chunk 1" + assert reasoning_contents[1] == "Thinking chunk 2" + assert reasoning_contents[2] == "Thinking chunk 3" # This was previously None + assert reasoning_contents[3] == "Thinking chunk 4" # This was previously None + assert reasoning_contents[4] == "Thinking chunk 5" # This was previously None + + # Verify none of them are None + for i, rc in enumerate(reasoning_contents): + assert rc is not None, f"Chunk {i+1} reasoning_content should not be None" + + def test_thinking_to_content_transition(self): + """ + Test that transition from thinking to regular content works correctly. + """ + iterator = OllamaChatCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + ) + + # First: thinking chunks + thinking_chunk = { + "model": "deepseek-r1", + "message": {"role": "assistant", "thinking": "Let me think about this..."}, + "done": False, + } + result1 = iterator.chunk_parser(thinking_chunk) + assert result1.choices[0].delta.reasoning_content == "Let me think about this..." + assert result1.choices[0].delta.content is None + + # Then: regular content chunk + content_chunk = { + "model": "deepseek-r1", + "message": {"role": "assistant", "content": "Here is my answer."}, + "done": False, + } + result2 = iterator.chunk_parser(content_chunk) + assert result2.choices[0].delta.content == "Here is my answer." + # reasoning_content is not set when there's no thinking in the chunk + assert getattr(result2.choices[0].delta, 'reasoning_content', None) is None + + def test_think_tags_in_content(self): + """ + Test that tags embedded in content are properly parsed. + """ + iterator = OllamaChatCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + ) + + # Content with tag + chunk1 = { + "model": "deepseek-r1", + "message": {"role": "assistant", "content": "I need to analyze this"}, + "done": False, + } + result1 = iterator.chunk_parser(chunk1) + assert result1.choices[0].delta.reasoning_content == "I need to analyze this" + assert result1.choices[0].delta.content is None + + # Content with tag (end of thinking) + chunk2 = { + "model": "deepseek-r1", + "message": {"role": "assistant", "content": "The answer is 42."}, + "done": False, + } + result2 = iterator.chunk_parser(chunk2) + assert result2.choices[0].delta.content == "The answer is 42." + # reasoning_content is not set when it's regular content + assert getattr(result2.choices[0].delta, 'reasoning_content', None) is None + + def test_done_chunk_with_thinking(self): + """ + Test that the final chunk with done=True and thinking content works. + """ + iterator = OllamaChatCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + ) + + # Final chunk with thinking + done_chunk = { + "model": "deepseek-r1", + "message": {"role": "assistant", "thinking": "Final thought"}, + "done": True, + "done_reason": "stop", + } + result = iterator.chunk_parser(done_chunk) + assert result.choices[0].delta.reasoning_content == "Final thought" + assert result.choices[0].finish_reason == "stop" + + diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 5f087363797..c0695bf3588 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -2,9 +2,10 @@ Tests for OpenAI GPT transformation (litellm/llms/openai/chat/gpt_transformation.py) """ -import pytest -import sys import os +import sys + +import pytest sys.path.insert(0, os.path.abspath("../../../../..")) @@ -73,6 +74,17 @@ class TestOpenAIGPTConfig: for param in base_expected_params: assert param in supported_params, f"Expected '{param}' in supported params" + def test_prompt_cache_key_supported(self): + """Test that 'prompt_cache_key' is in supported params for OpenAI chat completion models. + + OpenAI's Chat Completions API supports prompt_cache_key for cache routing optimization. + """ + supported_params = self.config.get_supported_openai_params("gpt-4.1-nano") + assert "prompt_cache_key" in supported_params + + supported_params = self.config.get_supported_openai_params("gpt-4.1") + assert "prompt_cache_key" in supported_params + class TestGetOptionalParamsIntegration: """Integration tests using litellm.get_optional_params()""" @@ -123,3 +135,14 @@ class TestGetOptionalParamsIntegration: # Both should include user assert regular_params.get("user") == "my-end-user" assert responses_params.get("user") == "my-end-user" + + def test_prompt_cache_key_in_optional_params(self): + """Test that 'prompt_cache_key' flows through get_optional_params for OpenAI models.""" + from litellm.utils import get_optional_params + + optional_params = get_optional_params( + model="gpt-4.1-nano", + custom_llm_provider="openai", + prompt_cache_key="test-cache-key-123", + ) + assert optional_params.get("prompt_cache_key") == "test-cache-key-123" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index a9c27e30930..c474461e0a2 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -7,6 +7,7 @@ from litellm.llms.vertex_ai.gemini.transformation import ( check_if_part_exists_in_parts, ) from litellm.types.llms.vertex_ai import BlobType +from litellm.types.utils import Message def test_check_if_part_exists_in_parts(): @@ -735,8 +736,9 @@ def test_file_data_field_order(): Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order. """ import json + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - + # Test with HTTPS URL and explicit format (audio file) file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123" format = "audio/mpeg" @@ -770,8 +772,9 @@ def test_file_data_field_order(): def test_file_data_field_order_gcs_urls(): """Test that GCS URLs also maintain correct field order.""" import json + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - + # Test with GCS URL gcs_url = "gs://bucket/audio.mp3" @@ -802,11 +805,14 @@ def test_extract_file_data_with_path_object(): Related issue: Files uploaded with wrong MIME type cause Gemini API to reject requests where the specified format doesn't match the uploaded file's MIME type. """ - from pathlib import Path - from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data - import tempfile import os - + import tempfile + from pathlib import Path + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + # Create a temporary MP3 file with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp: tmp.write(b"fake mp3 content") @@ -835,10 +841,13 @@ def test_extract_file_data_with_path_object(): def test_extract_file_data_with_string_path(): """Test that filename is correctly extracted from string paths.""" - from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data - import tempfile import os - + import tempfile + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + # Create a temporary WAV file with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: tmp.write(b"fake wav content") @@ -866,8 +875,10 @@ def test_extract_file_data_with_string_path(): def test_extract_file_data_with_tuple_format(): """Test that tuple format (with explicit content_type) still works correctly.""" - from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data - + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + # Test with tuple format: (filename, content, content_type) filename = "test_audio.mp3" content = b"test audio content" @@ -883,10 +894,13 @@ def test_extract_file_data_with_tuple_format(): def test_extract_file_data_fallback_to_octet_stream(): """Test that unknown file types fall back to application/octet-stream.""" - from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data - import tempfile import os - + import tempfile + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + # Create a temporary file with unknown extension with tempfile.NamedTemporaryFile(suffix=".xyz123", delete=False) as tmp: tmp.write(b"unknown content") @@ -1086,3 +1100,227 @@ def test_convert_tool_response_with_nested_file_object(): assert "mime_type" in inline_data assert inline_data["mime_type"] == "application/pdf" assert inline_data["data"] == test_pdf_base64 + +def test_assistant_message_with_images_field(): + """ + Test that assistant messages with images field are properly converted to Gemini format. + + This handles the case where an assistant message contains generated images in the + `images` field (e.g., from image generation models like gemini-2.5-flash-image). + The images should be converted to inline_data parts in the Gemini format. + """ + # Create a small test image (1x1 red pixel PNG) + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + # Create messages with assistant message containing images field + messages = [ + { + "role": "user", + "content": "Generate an image of a banana wearing a costume that says LiteLLM" + }, + { + "role": "assistant", + "content": "Here's your banana in a LiteLLM costume!", + "images": [ + { + "image_url": { + "url": image_data_uri, + "detail": "auto" + }, + "index": 0, + "type": "image_url" + } + ] + } + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify structure + assert len(contents) == 2, f"Expected 2 content blocks, got {len(contents)}" + + # Verify user message + assert contents[0]["role"] == "user" + assert len(contents[0]["parts"]) == 1 + assert contents[0]["parts"][0]["text"] == "Generate an image of a banana wearing a costume that says LiteLLM" + + # Verify assistant message + assert contents[1]["role"] == "model" + assert len(contents[1]["parts"]) == 2, f"Expected 2 parts (text + image), got {len(contents[1]['parts'])}" + + # Find text part and inline_data part + text_part = None + inline_data_part = None + for part in contents[1]["parts"]: + if "text" in part: + text_part = part + elif "inline_data" in part: + inline_data_part = part + + # Verify text part + assert text_part is not None, "Missing text part in assistant message" + assert text_part["text"] == "Here's your banana in a LiteLLM costume!" + + # Verify inline_data part (image) + assert inline_data_part is not None, "Missing inline_data part in assistant message" + inline_data: BlobType = inline_data_part["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "image/png" + assert inline_data["data"] == test_image_base64 + + +def test_assistant_message_with_multiple_images(): + """Test that assistant messages with multiple images are properly converted.""" + # Create two test images + test_image1_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + test_image2_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" + image1_data_uri = f"data:image/png;base64,{test_image1_base64}" + image2_data_uri = f"data:image/jpeg;base64,{test_image2_base64}" + + messages = [ + { + "role": "user", + "content": "Generate two images" + }, + { + "role": "assistant", + "content": "Here are your images:", + "images": [ + { + "image_url": { + "url": image1_data_uri, + "detail": "auto" + }, + "index": 0, + "type": "image_url" + }, + { + "image_url": { + "url": image2_data_uri, + "detail": "high" + }, + "index": 1, + "type": "image_url" + } + ] + } + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify assistant message has 3 parts (1 text + 2 images) + assert contents[1]["role"] == "model" + assert len(contents[1]["parts"]) == 3, f"Expected 3 parts (text + 2 images), got {len(contents[1]['parts'])}" + + # Count inline_data parts + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert len(inline_data_parts) == 2, f"Expected 2 inline_data parts, got {len(inline_data_parts)}" + + # Verify first image + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" + assert inline_data_parts[0]["inline_data"]["data"] == test_image1_base64 + + # Verify second image + assert inline_data_parts[1]["inline_data"]["mime_type"] == "image/jpeg" + assert inline_data_parts[1]["inline_data"]["data"] == test_image2_base64 + + +def test_assistant_message_with_images_using_message_object(): + """Test that Message objects with images field are properly converted.""" + # Create a small test image + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + # Create messages using Message object (as returned by LiteLLM) + user_message = { + "role": "user", + "content": "Generate an image" + } + + assistant_message = Message( + content="Here's your image!", + role="assistant", + tool_calls=None, + function_call=None, + images=[ + { + "image_url": { + "url": image_data_uri, + "detail": "auto" + }, + "index": 0, + "type": "image_url" + } + ] + ) + + messages = [user_message, assistant_message] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify assistant message has both text and image + assert contents[1]["role"] == "model" + assert len(contents[1]["parts"]) == 2 + + # Verify image was converted + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert len(inline_data_parts) == 1 + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" + assert inline_data_parts[0]["inline_data"]["data"] == test_image_base64 + + +def test_assistant_message_with_images_in_conversation_history(): + """ + Test multi-turn conversation where assistant message with images is in history. + + This simulates the real use case where: + 1. User asks for image generation + 2. Assistant generates image (with images field) + 3. User asks follow-up question about the image + """ + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + messages = [ + { + "role": "user", + "content": "Generate an image of a cat" + }, + { + "role": "assistant", + "content": "Here's a cat image:", + "images": [ + { + "image_url": { + "url": image_data_uri, + "detail": "auto" + }, + "index": 0, + "type": "image_url" + } + ] + }, + { + "role": "user", + "content": "Can you make it more colorful?" + } + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify structure: user -> model (with image) -> user + assert len(contents) == 3 + assert contents[0]["role"] == "user" + assert contents[1]["role"] == "model" + assert contents[2]["role"] == "user" + + # Verify assistant message has image in history + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert len(inline_data_parts) == 1 + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" \ No newline at end of file diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 4bcafd4c57e..3a49880ff16 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -459,7 +459,7 @@ def test_vertex_ai_partner_models_anthropic_remove_prompt_caching_scope_beta_hea assert PROMPT_CACHING_BETA_HEADER not in ( beta_header or "" ), f"{PROMPT_CACHING_BETA_HEADER} should be filtered out" - assert "other-feature" in ( + assert "other-feature" not in ( beta_header or "" ), "Other non-excluded beta headers should remain" assert "web-search-2025-03-05" in ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index e4b4d3ba189..d2b00d61d1b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -2,6 +2,8 @@ import pytest from unittest.mock import AsyncMock, MagicMock, patch +from fastapi import HTTPException + # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control @@ -86,6 +88,70 @@ async def test_authorize_endpoint_includes_response_type(): assert "scope=read+write" in response.headers["location"] +@pytest.mark.asyncio +async def test_authorize_endpoint_preserves_existing_query_params(): + """Test that authorize endpoint merges OAuth params with existing query params in authorization_url""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + + # Authorization URL already has query params (e.g. multi-tenant OAuth) + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize?tenant=system", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "mocked_encrypted_state" + + response = await authorize( + request=mock_request, + client_id="test_client_id", + mcp_server_name="test_oauth", + redirect_uri="https://client.example.com/callback", + state="test_state", + ) + + location = response.headers["location"] + + # Must NOT have double '?' — existing params must be merged correctly + assert location.count("?") == 1, ( + f"Expected exactly one '?' in URL but got {location.count('?')}: {location}" + ) + assert "tenant=system" in location + assert "client_id=test_client_id" in location + assert "response_type=code" in location + assert "scope=read+write" in location + + @pytest.mark.asyncio async def test_authorize_endpoint_forwards_pkce_parameters(): """Test that authorize endpoint forwards PKCE parameters (code_challenge and code_challenge_method)""" @@ -260,10 +326,16 @@ async def test_register_client_without_mcp_server_name_returns_dummy(): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) from fastapi import Request except ImportError: pytest.skip("MCP discoverable endpoints not available") + # Clear registry to ensure no OAuth2 servers exist (otherwise resolver would find one) + global_mcp_server_manager.registry.clear() + mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.litellm.example/" mock_request.headers = {} @@ -680,10 +752,16 @@ async def test_register_client_respects_x_forwarded_proto(): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) from fastapi import Request except ImportError: pytest.skip("MCP discoverable endpoints not available") + # Clear registry to ensure no OAuth2 servers exist (otherwise resolver would find one) + global_mcp_server_manager.registry.clear() + # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) mock_request.base_url = "http://proxy.litellm.example/" # HTTP @@ -1017,3 +1095,263 @@ def test_get_request_base_url_comprehensive( f"X-Forwarded-Host={x_forwarded_host}, " f"X-Forwarded-Port={x_forwarded_port}" ) + + +# ------------------------------------------------------------------- +# Tests for root-level OAuth endpoint resolution (no server name) +# ------------------------------------------------------------------- + + +def _create_oauth2_server( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + client_id="test_client_id", + client_secret="test_client_secret", +): + """Helper to create a mock OAuth2 MCPServer.""" + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + + return MCPServer( + server_id=server_id, + name=name, + server_name=server_name, + alias=alias, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=client_id, + client_secret=client_secret, + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + + +@pytest.mark.asyncio +async def test_authorize_root_resolves_single_oauth2_server(): + """When /authorize is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "mocked_encrypted_state" + + # Call /authorize WITHOUT mcp_server_name, with dummy_client as client_id + response = await authorize( + request=mock_request, + client_id="dummy_client", + mcp_server_name=None, + redirect_uri="http://localhost:62646/callback", + state="test_state", + ) + + # Should resolve to the single OAuth2 server and redirect + assert response.status_code == 307 + location = response.headers["location"] + assert "https://provider.com/oauth/authorize" in location + assert "client_id=test_client_id" in location + finally: + global_mcp_server_manager.registry.clear() + + +@pytest.mark.asyncio +async def test_authorize_root_fails_with_multiple_oauth2_servers(): + """When /authorize is hit without server name and multiple OAuth2 servers exist, return 404.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + server1 = _create_oauth2_server( + server_id="server1", name="server1", server_name="server1", alias="server1" + ) + server2 = _create_oauth2_server( + server_id="server2", name="server2", server_name="server2", alias="server2" + ) + global_mcp_server_manager.registry[server1.server_id] = server1 + global_mcp_server_manager.registry[server2.server_id] = server2 + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with pytest.raises(HTTPException) as exc_info: + await authorize( + request=mock_request, + client_id="dummy_client", + mcp_server_name=None, + redirect_uri="http://localhost:62646/callback", + state="test_state", + ) + assert exc_info.value.status_code == 404 + assert "MCP server not found" in str(exc_info.value.detail) + finally: + global_mcp_server_manager.registry.clear() + + +@pytest.mark.asyncio +async def test_token_root_resolves_single_oauth2_server(): + """When /token is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + token_endpoint, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + mock_response = MagicMock() + mock_response.json.return_value = { + "access_token": "ya29.test_token", + "token_type": "Bearer", + "expires_in": 3599, + } + mock_response.raise_for_status = MagicMock() + + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + try: + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = mock_async_client + + # Call /token WITHOUT mcp_server_name + response = await token_endpoint( + request=mock_request, + grant_type="authorization_code", + code="test_auth_code", + redirect_uri="http://localhost:62646/callback", + client_id="dummy_client", + mcp_server_name=None, + client_secret=None, + code_verifier="test_verifier", + ) + + # Should resolve and exchange token with the upstream server + import json + + token_data = json.loads(response.body) + assert token_data["access_token"] == "ya29.test_token" + + # Verify it called the correct upstream token URL + call_args = mock_async_client.post.call_args + assert call_args.args[0] == "https://provider.com/oauth/token" + finally: + global_mcp_server_manager.registry.clear() + + +@pytest.mark.asyncio +async def test_register_root_resolves_single_oauth2_server(): + """When /register is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value={}), + ): + result = await register_client(request=mock_request, mcp_server_name=None) + + # Should resolve to the single server and return its name as client_id + assert result["client_id"] == "test_oauth" + assert "redirect_uris" in result + finally: + global_mcp_server_manager.registry.clear() + + +@pytest.mark.asyncio +async def test_discovery_root_includes_server_name_prefix(): + """When root discovery is hit and exactly 1 OAuth2 server exists, include server name in URLs.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_authorization_server_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from fastapi import Request + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + # Call with mcp_server_name=None (root discovery) + response = _build_oauth_authorization_server_response( + request=mock_request, + mcp_server_name=None, + ) + + # Should resolve to the single server and include its name in endpoint URLs + assert "/test_oauth/authorize" in response["authorization_endpoint"] + assert "/test_oauth/token" in response["token_endpoint"] + assert "/test_oauth/register" in response["registration_endpoint"] + assert response["scopes_supported"] == ["read", "write"] + finally: + global_mcp_server_manager.registry.clear() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index abb8dd49159..b4b5811666b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1043,7 +1043,7 @@ class TestMCPServerManager: "litellm.proxy._experimental.mcp_server.tool_registry.global_mcp_tool_registry.register_tool", return_value=None, ): - manager._register_openapi_tools( + await manager._register_openapi_tools( spec_path=str(spec_path), server=server, base_url="https://example.com", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 077948cbeb9..4f93270c162 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -157,6 +157,160 @@ class TestExecuteWithMcpClient: } + @pytest.mark.asyncio + async def test_m2m_credentials_forwarded_to_server_model(self, monkeypatch): + """M2M OAuth credentials (client_id, client_secret) from the nested + ``credentials`` dict must be forwarded to the MCPServer model so that + ``has_client_credentials`` returns True and the proxy auto-fetches tokens.""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured["server"] = kwargs.get("server") + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="m2m-server", + url="https://example.com", + auth_type=MCPAuth.oauth2, + token_url="https://auth.example.com/token", + credentials={ + "client_id": "my-id", + "client_secret": "my-secret", + "scopes": ["read", "write"], + }, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, ok_operation + ) + + assert result["status"] == "ok" + server = captured["server"] + assert server.client_id == "my-id" + assert server.client_secret == "my-secret" + assert server.token_url == "https://auth.example.com/token" + assert server.scopes == ["read", "write"] + assert server.has_client_credentials is True + + @pytest.mark.asyncio + async def test_m2m_drops_incoming_oauth2_headers(self, monkeypatch): + """For M2M OAuth servers the incoming Authorization header (which carries + the litellm API key) must NOT be forwarded as extra_headers — otherwise + it overwrites the auto-fetched M2M token.""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured["extra_headers"] = kwargs.get("extra_headers") + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="m2m-server", + url="https://example.com", + auth_type=MCPAuth.oauth2, + token_url="https://auth.example.com/token", + credentials={ + "client_id": "my-id", + "client_secret": "my-secret", + }, + ) + + incoming_oauth2 = {"Authorization": "Bearer sk-litellm-api-key"} + result = await rest_endpoints._execute_with_mcp_client( + payload, + ok_operation, + oauth2_headers=incoming_oauth2, + ) + + assert result["status"] == "ok" + # The incoming Authorization must be dropped — extra_headers should + # contain no oauth2 headers (only static_headers, which are None here). + assert captured["extra_headers"] is None or "Authorization" not in captured["extra_headers"] + + @pytest.mark.asyncio + async def test_catches_exception_group(self, monkeypatch): + """MCP SDK's anyio TaskGroup raises BaseExceptionGroup which does not + inherit from Exception. The handler must catch it and return an error + dict instead of letting a raw 500 propagate.""" + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + raise BaseExceptionGroup( + "test group", [RuntimeError("Cancelled via cancel scope")] + ) + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="bad-server", + url="https://example.com", + auth_type=MCPAuth.none, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, ok_operation + ) + + assert result["status"] == "error" + assert result["error"] is True + assert "Failed to connect to MCP server" in result["message"] + # Error message must not leak raw exception details + assert "cancel scope" not in result["message"] + + class TestTestConnection: def test_requires_auth_dependency(self): route = _get_route("/mcp-rest/test/connection", "POST") diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index a745ac3de13..65bb329e7ea 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -108,6 +108,22 @@ def test_virtual_key_allowed_routes_with_litellm_routes_member_name_allowed(): assert result is True +def test_virtual_key_mcp_routes_allows_v1_mcp_server(): + """Regression test for #20325: allow virtual keys to list MCP servers.""" + + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["mcp_routes"], + ) + + result = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/v1/mcp/server", + valid_token=valid_token, + ) + + assert result is True + + def test_virtual_key_allowed_routes_with_litellm_routes_member_name_denied(): """Test that virtual key is denied when route is not in the allowed LiteLLMRoutes group""" @@ -199,14 +215,14 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route): def test_google_routes_with_dynamic_model_names_recognized_as_llm_api_route(route): """ Test that Google routes with dynamic model names (including custom names) are recognized as LLM API routes. - + This test verifies the fix for the issue where routes like: /v1beta/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent were incorrectly classified as "custom admin only route" instead of LLM API routes. - + The fix adds pattern matching for Google routes with placeholders like {model_name}. """ - + # Test that the route is recognized as an LLM API route assert RouteChecks.is_llm_api_route(route) is True @@ -214,28 +230,28 @@ def test_google_routes_with_dynamic_model_names_recognized_as_llm_api_route(rout def test_google_routes_with_dynamic_model_names_accessible_to_internal_users(): """ Test that internal users can access Google routes with dynamic model names. - + This ensures that routes like /v1beta/models/{model_name}:generateContent are properly accessible to internal users and not blocked as admin-only routes. """ - + # Create an internal user object user_obj = LiteLLM_UserTable( user_id="test_user", user_email="test@example.com", user_role=LitellmUserRoles.INTERNAL_USER.value, ) - + # Create an internal user API key auth valid_token = UserAPIKeyAuth( user_id="test_user", user_role=LitellmUserRoles.INTERNAL_USER.value, ) - + # Create a mock request request = MagicMock(spec=Request) request.query_params = {} - + # Test that calling Google route with dynamic model name does NOT raise an exception try: RouteChecks.non_proxy_admin_allowed_routes_check( @@ -263,11 +279,13 @@ def test_virtual_key_allowed_routes_with_multiple_litellm_routes_member_names(): # Test that routes from both groups are allowed result1 = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/chat/completions", valid_token=valid_token # This is in openai_routes + route="/chat/completions", + valid_token=valid_token, # This is in openai_routes ) result2 = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/user/info", valid_token=valid_token # This is in info_routes + route="/user/info", + valid_token=valid_token, # This is in info_routes ) assert result1 is True @@ -288,11 +306,13 @@ def test_virtual_key_allowed_routes_with_mixed_member_names_and_explicit_routes( # Test that both info routes and explicit custom route are allowed result1 = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/user/info", valid_token=valid_token # This is in info_routes + route="/user/info", + valid_token=valid_token, # This is in info_routes ) result2 = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/custom/route", valid_token=valid_token # This is explicitly listed + route="/custom/route", + valid_token=valid_token, # This is explicitly listed ) assert result1 is True @@ -323,7 +343,8 @@ def test_virtual_key_allowed_routes_with_no_member_names_only_explicit(): # Test that non-allowed route raises HTTPException with pytest.raises(HTTPException) as exc_info: RouteChecks.is_virtual_key_allowed_to_call_route( - route="/user/info", valid_token=valid_token # Not in allowed routes + route="/user/info", + valid_token=valid_token, # Not in allowed routes ) assert exc_info.value.status_code == 403 @@ -372,12 +393,15 @@ def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints(): }, } - with patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", - mock_registered_routes, - ), patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path", - return_value="/", + with ( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + mock_registered_routes, + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path", + return_value="/", + ), ): # Create a virtual key with llm_api_routes permission valid_token = UserAPIKeyAuth( @@ -421,12 +445,15 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): }, } - with patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", - mock_registered_routes, - ), patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path", - return_value="/", + with ( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + mock_registered_routes, + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path", + return_value="/", + ), ): # Create a virtual key without llm_api_routes permission valid_token = UserAPIKeyAuth( @@ -442,7 +469,9 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): ) assert exc_info.value.status_code == 403 - assert "Virtual key is not allowed to call this route" in str(exc_info.value.detail) + assert "Virtual key is not allowed to call this route" in str( + exc_info.value.detail + ) def test_check_passthrough_route_access_key_metadata_exact_match(): @@ -833,6 +862,7 @@ def test_videos_route_with_virtual_key_llm_api_routes(): result is True ), f"Virtual key with llm_api_routes should be able to access {route}" + def test_non_proxy_admin_wildcard_allowed_routes(): """Test that nonproxy admin users can still use wildcard routes""" @@ -847,7 +877,7 @@ def test_non_proxy_admin_wildcard_allowed_routes(): user_role=LitellmUserRoles.INTERNAL_USER.value, allowed_routes=["/scim/*"], ) - + request = MagicMock(spec=Request) request.query_params = {} @@ -864,14 +894,14 @@ def test_non_proxy_admin_wildcard_allowed_routes(): def test_proxy_admin_viewer_can_access_global_spend_tags(): """ Test that proxy_admin_viewer can access /global/spend/tags endpoint. - + This test verifies the fix for the issue where proxy_admin_viewer was getting 403 errors when trying to access /global/spend/tags endpoint. - + Related: Slack thread from 10/9/2025 - Erik Kristensen reported this issue. proxy_admin_viewer role should have access to "view all spend" endpoints. """ - + # Create a proxy admin viewer user object user_obj = LiteLLM_UserTable( user_id="viewer_user", @@ -912,8 +942,12 @@ def test_route_in_additional_public_routes_wildcard_match(): """ from litellm.proxy.auth.auth_utils import route_in_additonal_public_routes - with patch("litellm.proxy.proxy_server.general_settings", {"public_routes": ["/api/*"]}), \ - patch("litellm.proxy.proxy_server.premium_user", True): + with ( + patch( + "litellm.proxy.proxy_server.general_settings", {"public_routes": ["/api/*"]} + ), + patch("litellm.proxy.proxy_server.premium_user", True), + ): # Wildcard should match subpaths assert route_in_additonal_public_routes("/api/users") is True assert route_in_additonal_public_routes("/api/users/123") is True @@ -927,11 +961,15 @@ def test_route_in_additional_public_routes_exact_match(): """ from litellm.proxy.auth.auth_utils import route_in_additonal_public_routes - with patch("litellm.proxy.proxy_server.general_settings", {"public_routes": ["/health", "/status"]}), \ - patch("litellm.proxy.proxy_server.premium_user", True): + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"public_routes": ["/health", "/status"]}, + ), + patch("litellm.proxy.proxy_server.premium_user", True), + ): # Exact matches should work assert route_in_additonal_public_routes("/health") is True assert route_in_additonal_public_routes("/status") is True # Non-matching routes should fail assert route_in_additonal_public_routes("/other") is False - diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 1f379f4371e..9b7b7f46155 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -41,6 +41,12 @@ def test_get_api_key(): ("Basic sk-12345678", "sk-12345678", "Basic sk-12345678"), ("bearer sk-12345678", "sk-12345678", "bearer sk-12345678"), ("sk-12345678", "sk-12345678", "sk-12345678"), + # AWS Signature V4 format (LangChain AWS SDK) + ( + "AWS4-HMAC-SHA256 Credential=Bearer sk-12345678/20260210/us-east-1/bedrock/aws4_request, SignedHeaders=host, Signature=abc123", + "sk-12345678", + "AWS4-HMAC-SHA256 Credential=Bearer sk-12345678/20260210/us-east-1/bedrock/aws4_request, SignedHeaders=host, Signature=abc123", + ), ], ) def test_get_api_key_with_custom_litellm_key_header( @@ -243,10 +249,10 @@ async def test_proxy_admin_expired_key_from_cache(): Regression test for issue where PROXY_ADMIN keys from cache skipped expiration check. """ from datetime import datetime, timedelta, timezone - + from fastapi import Request from starlette.datastructures import URL - + from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, @@ -255,7 +261,7 @@ async def test_proxy_admin_expired_key_from_cache(): ) from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder from litellm.proxy.proxy_server import hash_token - + # Create an expired PROXY_ADMIN key api_key = "sk-test-proxy-admin-key" hashed_key = hash_token(api_key) @@ -368,7 +374,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget(): from user_obj attributes. """ from datetime import datetime - + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index 88d31e993dd..9d0c771e1d9 100644 --- a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -31,6 +31,7 @@ def test_ui_discovery_endpoints_with_defaults(): assert data["proxy_base_url"] is None assert data["auto_redirect_to_sso"] is False assert data["admin_ui_disabled"] is False + assert data["sso_configured"] is False def test_ui_discovery_endpoints_with_custom_server_root_path(): @@ -50,6 +51,7 @@ def test_ui_discovery_endpoints_with_custom_server_root_path(): assert data["server_root_path"] == "/litellm" assert data["proxy_base_url"] is None assert data["auto_redirect_to_sso"] is False + assert data["sso_configured"] is False def test_ui_discovery_endpoints_with_proxy_base_url_when_set(): @@ -69,6 +71,7 @@ def test_ui_discovery_endpoints_with_proxy_base_url_when_set(): assert data["server_root_path"] == "/" assert data["proxy_base_url"] == "https://proxy.example.com" assert data["auto_redirect_to_sso"] is False + assert data["sso_configured"] is False def test_ui_discovery_endpoints_with_sso_configured_and_auto_redirect_enabled(): @@ -88,6 +91,30 @@ def test_ui_discovery_endpoints_with_sso_configured_and_auto_redirect_enabled(): assert data["server_root_path"] == "/litellm" assert data["proxy_base_url"] == "https://proxy.example.com" assert data["auto_redirect_to_sso"] is True + assert data["sso_configured"] is True + + +def test_ui_discovery_endpoints_with_sso_configured_and_auto_redirect_not_set_defaults_to_false(): + """When SSO is configured but AUTO_REDIRECT_UI_LOGIN_TO_SSO is not set, defaults to False.""" + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com"), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), \ + patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False): + # Ensure AUTO_REDIRECT_UI_LOGIN_TO_SSO is not set (simulate default) + os.environ.pop("AUTO_REDIRECT_UI_LOGIN_TO_SSO", None) + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["server_root_path"] == "/litellm" + assert data["proxy_base_url"] == "https://proxy.example.com" + assert data["auto_redirect_to_sso"] is False + assert data["sso_configured"] is True def test_ui_discovery_endpoints_with_sso_configured_but_auto_redirect_disabled(): @@ -107,6 +134,7 @@ def test_ui_discovery_endpoints_with_sso_configured_but_auto_redirect_disabled() assert data["server_root_path"] == "/litellm" assert data["proxy_base_url"] == "https://proxy.example.com" assert data["auto_redirect_to_sso"] is False + assert data["sso_configured"] is True def test_ui_discovery_endpoints_with_sso_not_configured_but_auto_redirect_enabled(): @@ -126,6 +154,7 @@ def test_ui_discovery_endpoints_with_sso_not_configured_but_auto_redirect_enable assert data["server_root_path"] == "/" assert data["proxy_base_url"] is None assert data["auto_redirect_to_sso"] is False + assert data["sso_configured"] is False def test_ui_discovery_endpoints_both_routes_return_same_data(): @@ -164,6 +193,7 @@ def test_ui_discovery_endpoints_with_admin_ui_disabled(): assert data["proxy_base_url"] is None assert data["auto_redirect_to_sso"] is False assert data["admin_ui_disabled"] is True + assert data["sso_configured"] is False def test_ui_discovery_endpoints_with_admin_ui_enabled(): @@ -184,4 +214,5 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled(): assert data["proxy_base_url"] is None assert data["auto_redirect_to_sso"] is False assert data["admin_ui_disabled"] is False + assert data["sso_configured"] is False diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 5ec9b13408e..f01c23f7116 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -6,6 +6,7 @@ Tests PII detection and masking for different message formats import asyncio import os import sys +from contextlib import asynccontextmanager from unittest.mock import MagicMock, patch import pytest @@ -18,10 +19,41 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.presidio import ( _OPTIONAL_PresidioPIIMasking, ) +from litellm.exceptions import GuardrailRaisedException from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType from litellm.types.utils import Choices, Message, ModelResponse +def _make_mock_session_iterator(json_response): + """Create a mock _get_session_iterator that yields a session returning json_response.""" + + @asynccontextmanager + async def mock_iterator(): + class MockResponse: + async def json(self): + return json_response + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + class MockSession: + def post(self, *args, **kwargs): + return MockResponse() + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + yield MockSession() + + return mock_iterator + + @pytest.fixture def presidio_guardrail(): """Create a Presidio guardrail instance for testing""" @@ -889,39 +921,134 @@ async def test_analyze_text_error_dict_handling(): output_parse_pii=False, ) - # Mock the HTTP response to return error dict - class MockResponse: - async def json(self): - return {"error": "No text provided"} - - async def __aenter__(self): - return self - - async def __aexit__(self, *args): - pass - - class MockSession: - def post(self, *args, **kwargs): - return MockResponse() - - async def __aenter__(self): - return self - - async def __aexit__(self, *args): - pass - - with patch("aiohttp.ClientSession", return_value=MockSession()): + with patch.object( + presidio, + "_get_session_iterator", + _make_mock_session_iterator({"error": "No text provided"}), + ): result = await presidio.analyze_text( text="some text", presidio_config=None, request_data={}, ) - # Should return empty list when error dict is received - assert result == [], "Error dict should be handled gracefully" + assert result == [], "Error dict should be handled gracefully" print("✓ analyze_text error dict handling test passed") +@pytest.mark.asyncio +async def test_analyze_text_string_response_handling(): + """ + Test that analyze_text handles string responses from Presidio API. + + When Presidio returns a string (e.g. error message from websearch/hosted models), + should handle gracefully instead of crashing with TypeError about mapping vs str. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://mock-presidio:5002/", + presidio_anonymizer_api_base="http://mock-presidio:5001/", + output_parse_pii=False, + ) + + with patch.object( + presidio, + "_get_session_iterator", + _make_mock_session_iterator("Internal Server Error"), + ): + result = await presidio.analyze_text( + text="some text", + presidio_config=None, + request_data={}, + ) + assert result == [], "String response should be handled gracefully" + + +@pytest.mark.asyncio +async def test_analyze_text_invalid_response_raises_when_block_configured(): + """ + When pii_entities_config has BLOCK and Presidio returns invalid response, + should raise GuardrailRaisedException (fail-closed) rather than silently allowing content. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://mock-presidio:5002/", + presidio_anonymizer_api_base="http://mock-presidio:5001/", + output_parse_pii=False, + pii_entities_config={PiiEntityType.CREDIT_CARD: PiiAction.BLOCK}, + ) + + with patch.object( + presidio, + "_get_session_iterator", + _make_mock_session_iterator("Internal Server Error"), + ): + with pytest.raises(GuardrailRaisedException) as exc_info: + await presidio.analyze_text( + text="some text", + presidio_config=None, + request_data={}, + ) + assert "BLOCK" in str(exc_info.value) or "Presidio" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_analyze_text_invalid_response_raises_when_mask_configured(): + """ + When pii_entities_config has MASK and Presidio returns invalid response, + should raise GuardrailRaisedException (fail-closed) because PII masking is expected. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://mock-presidio:5002/", + presidio_anonymizer_api_base="http://mock-presidio:5001/", + output_parse_pii=False, + pii_entities_config={PiiEntityType.CREDIT_CARD: PiiAction.MASK}, + ) + + with patch.object( + presidio, + "_get_session_iterator", + _make_mock_session_iterator("Internal Server Error"), + ): + with pytest.raises(GuardrailRaisedException) as exc_info: + await presidio.analyze_text( + text="some text", + presidio_config=None, + request_data={}, + ) + assert "PII protection is configured" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_analyze_text_list_with_non_dict_items(): + """ + Test that analyze_text skips non-dict items in the result list. + + When Presidio returns a list containing strings (malformed response), + should skip invalid items and return parsed valid ones. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://mock-presidio:5002/", + presidio_anonymizer_api_base="http://mock-presidio:5001/", + output_parse_pii=False, + ) + + json_response = [ + {"entity_type": "PERSON", "start": 0, "end": 5, "score": 0.9}, + "invalid_string_item", + {"entity_type": "EMAIL", "start": 10, "end": 25, "score": 0.85}, + ] + with patch.object( + presidio, "_get_session_iterator", _make_mock_session_iterator(json_response) + ): + result = await presidio.analyze_text( + text="some text", + presidio_config=None, + request_data={}, + ) + assert len(result) == 2, "Should parse 2 valid dict items and skip the string" + assert result[0].get("entity_type") == "PERSON" + assert result[1].get("entity_type") == "EMAIL" + + @pytest.mark.asyncio async def test_tool_calling_complete_scenario( presidio_guardrail, mock_user_api_key, mock_cache diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 88f56c24067..c0f16c8b953 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -149,6 +149,111 @@ async def test_list_guardrails_v2_with_db_and_config( assert isinstance(config_guardrail.litellm_params, BaseLitellmParams) +@pytest.mark.asyncio +async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker): + """Test that sensitive litellm_params are masked for DB guardrails in list response""" + db_guardrail_with_secrets = { + "guardrail_id": "secret-db-guardrail", + "guardrail_name": "DB Guardrail with Secrets", + "litellm_params": { + "guardrail": "azure/text_moderations", + "mode": "pre_call", + "api_key": "sk-1234567890abcdef", + "api_base": "https://api.secret.example.com", + }, + "guardrail_info": {"description": "Test guardrail"}, + "created_at": datetime.now(), + "updated_at": datetime.now(), + } + + mock_prisma_client = mocker.Mock() + mock_prisma_client.db = mocker.Mock() + mock_prisma_client.db.litellm_guardrailstable = mocker.Mock() + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[db_guardrail_with_secrets] + ) + + mock_in_memory_handler = mocker.Mock() + mock_in_memory_handler.list_in_memory_guardrails.return_value = [] + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + + response = await list_guardrails_v2() + + assert len(response.guardrails) == 1 + guardrail = response.guardrails[0] + litellm_params = guardrail.litellm_params + if isinstance(litellm_params, dict): + params = litellm_params + else: + params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params) + + # Sensitive keys (containing "key", "secret", "token", etc.) should be masked + assert params["api_key"] != "sk-1234567890abcdef" + assert "****" in str(params["api_key"]) + # Non-sensitive keys should remain unchanged + assert params["guardrail"] == "azure/text_moderations" + assert params["mode"] == "pre_call" + assert params["api_base"] == "https://api.secret.example.com" + + +@pytest.mark.asyncio +async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mocker): + """Test that sensitive litellm_params are masked for in-memory/config guardrails in list response""" + config_guardrail_with_secrets = { + "guardrail_id": "secret-config-guardrail", + "guardrail_name": "Config Guardrail with Secrets", + "litellm_params": { + "guardrail": "bedrock", + "mode": "during_call", + "api_key": "my-secret-bedrock-key", + "vertex_credentials": "{sensitive_creds}", + }, + "guardrail_info": {"description": "Test guardrail from config"}, + } + + mock_prisma_client = mocker.Mock() + mock_prisma_client.db = mocker.Mock() + mock_prisma_client.db.litellm_guardrailstable = mocker.Mock() + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[] + ) + + mock_in_memory_handler = mocker.Mock() + mock_in_memory_handler.list_in_memory_guardrails.return_value = [ + config_guardrail_with_secrets + ] + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + + response = await list_guardrails_v2() + + assert len(response.guardrails) == 1 + guardrail = response.guardrails[0] + litellm_params = guardrail.litellm_params + if isinstance(litellm_params, dict): + params = litellm_params + else: + params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params) + + # Sensitive keys should be masked + assert params["api_key"] != "my-secret-bedrock-key" + assert "****" in str(params["api_key"]) + assert params["vertex_credentials"] != "{sensitive_creds}" + assert "****" in str(params["vertex_credentials"]) + # Non-sensitive keys should remain unchanged + assert params["guardrail"] == "bedrock" + assert params["mode"] == "during_call" + + @pytest.mark.asyncio async def test_get_guardrail_info_from_db(mocker, mock_prisma_client): """Test getting guardrail info from DB""" diff --git a/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py b/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py new file mode 100644 index 00000000000..4a5d901b74d --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py @@ -0,0 +1,293 @@ +""" +Tests that guardrails (post_call_success_hook) fire for image generation requests. + +The /images/generations endpoint in proxy/image_endpoints/endpoints.py calls +proxy_logging_obj.post_call_success_hook after a successful image generation. +These tests verify: +1. CustomGuardrail.async_post_call_success_hook is invoked for image generation. +2. A guardrail can inspect and transform the image response. +3. A guardrail that raises blocks the response (exception propagates). +""" + +import os +import sys +from typing import Any, Optional +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ImageObject, ImageResponse + + +def _make_image_response(**kwargs) -> ImageResponse: + """Helper to build a minimal ImageResponse for tests.""" + return ImageResponse( + data=[ImageObject(url="https://example.com/img.png")], + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# 1. Hook is invoked for image generation responses +# --------------------------------------------------------------------------- + + +class TrackingGuardrail(CustomGuardrail): + """Guardrail that records whether it was called and with what args.""" + + def __init__(self): + super().__init__( + guardrail_name="tracking_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + self.called = False + self.received_data: Optional[dict] = None + self.received_response: Optional[Any] = None + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + self.called = True + self.received_data = data + self.received_response = response + return response + + +@pytest.mark.asyncio +async def test_post_call_success_hook_invoked_for_image_generation(): + """ + Verify that a default-on guardrail's async_post_call_success_hook is + called when ProxyLogging.post_call_success_hook is invoked with an + ImageResponse (the same path used by the /images/generations endpoint). + """ + guardrail = TrackingGuardrail() + image_response = _make_image_response() + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + data = {"model": "dall-e-3", "prompt": "A sunset over mountains"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.post_call_success_hook( + data=data, + response=image_response, + user_api_key_dict=user_api_key_dict, + ) + + assert guardrail.called is True, "Guardrail hook was not invoked for image generation" + assert guardrail.received_data is not None + assert guardrail.received_data["model"] == "dall-e-3" + assert isinstance(guardrail.received_response, ImageResponse) + # The response should be passed through unchanged + assert result is image_response + + +# --------------------------------------------------------------------------- +# 2. Guardrail can transform image generation response +# --------------------------------------------------------------------------- + + +class TransformingGuardrail(CustomGuardrail): + """Guardrail that replaces the image URL in the response.""" + + def __init__(self): + super().__init__( + guardrail_name="transforming_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + # Return a modified image response (e.g., watermarked URL) + return ImageResponse( + data=[ImageObject(url="https://example.com/watermarked.png")], + ) + + +@pytest.mark.asyncio +async def test_guardrail_can_transform_image_response(): + """ + Verify that a guardrail can replace the ImageResponse returned to the client. + """ + guardrail = TransformingGuardrail() + original_response = _make_image_response() + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + data = {"model": "dall-e-3", "prompt": "A sunset"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.post_call_success_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + assert result is not original_response + assert isinstance(result, ImageResponse) + assert result.data[0].url == "https://example.com/watermarked.png" + + +# --------------------------------------------------------------------------- +# 3. Guardrail that raises blocks the image response +# --------------------------------------------------------------------------- + + +class BlockingGuardrail(CustomGuardrail): + """Guardrail that raises on unsafe image prompts.""" + + def __init__(self): + super().__init__( + guardrail_name="blocking_guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + raise ValueError("Image content blocked by guardrail") + + +@pytest.mark.asyncio +async def test_guardrail_exception_propagates_for_image_generation(): + """ + Verify that an exception raised in a guardrail's post_call_success_hook + propagates up (the proxy endpoint wraps this in an error response). + """ + guardrail = BlockingGuardrail() + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + data = {"model": "dall-e-3", "prompt": "Something unsafe"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + with pytest.raises(ValueError, match="Image content blocked by guardrail"): + await proxy_logging.post_call_success_hook( + data=data, + response=_make_image_response(), + user_api_key_dict=user_api_key_dict, + ) + + +# --------------------------------------------------------------------------- +# 4. Non-guardrail CustomLogger also fires for image generation +# --------------------------------------------------------------------------- + + +class TrackingLogger(CustomLogger): + """Plain CustomLogger (not a guardrail) that tracks invocations.""" + + def __init__(self): + self.called = False + self.received_response = None + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + self.called = True + self.received_response = response + return response + + +@pytest.mark.asyncio +async def test_custom_logger_post_call_success_hook_fires_for_image_generation(): + """ + Verify that a plain CustomLogger (non-guardrail) callback also has its + async_post_call_success_hook invoked for image generation responses. + """ + logger = TrackingLogger() + image_response = _make_image_response() + + with patch("litellm.callbacks", [logger]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + data = {"model": "dall-e-3", "prompt": "A cat"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.post_call_success_hook( + data=data, + response=image_response, + user_api_key_dict=user_api_key_dict, + ) + + assert logger.called is True + assert isinstance(logger.received_response, ImageResponse) + assert result is image_response + + +# --------------------------------------------------------------------------- +# 5. Guardrail with should_run_guardrail=False is skipped +# --------------------------------------------------------------------------- + + +class OptInGuardrail(CustomGuardrail): + """Guardrail that is NOT default_on, so it only runs if explicitly requested.""" + + def __init__(self): + super().__init__( + guardrail_name="opt_in_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.post_call, + ) + self.called = False + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + self.called = True + return response + + +@pytest.mark.asyncio +async def test_non_default_guardrail_skipped_for_image_generation(): + """ + Verify that a guardrail with default_on=False is NOT invoked for image + generation unless the request explicitly enables it. + """ + guardrail = OptInGuardrail() + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + # No guardrails key in data -> should_run_guardrail returns False + data = {"model": "dall-e-3", "prompt": "A sunset"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + await proxy_logging.post_call_success_hook( + data=data, + response=_make_image_response(), + user_api_key_dict=user_api_key_dict, + ) + + assert guardrail.called is False, "Opt-in guardrail should not fire without explicit request" diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index d8c505223d9..b15b9d622e4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -21,6 +21,7 @@ sys.path.insert( def client_and_mocks(monkeypatch): # Setup MagicMock Prisma mock_prisma = MagicMock() + mock_table = MagicMock() mock_table.create = AsyncMock(side_effect=lambda *, data: data) mock_table.update = AsyncMock(side_effect=lambda *, where, data: {**where, **data}) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index b372476c3d6..8b7b5a6fb7a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -7,10 +7,28 @@ enterprise (premium) license check, but should still be applied so that users can intentionally clear previously-set fields. """ -from unittest.mock import patch +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch +import pytest + +from litellm.proxy._types import ( + Member, + LiteLLM_OrganizationMembershipTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, +) from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_admin, + _org_admin_can_invite_user, + _set_object_metadata_field, + _team_admin_can_invite_user, _update_metadata_fields, + _user_has_admin_privileges, + _user_has_admin_view, + admin_can_invite_user, ) @@ -160,3 +178,311 @@ class TestUpdateMetadataFieldsEmptyCollections: } _update_metadata_fields(updated_kv=updated_kv) mock_premium_check.assert_not_called() + + +class TestUserHasAdminView: + """Tests for _user_has_admin_view function.""" + + @pytest.mark.parametrize( + "user_role,expected", + [ + (LitellmUserRoles.PROXY_ADMIN, True), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, True), + (LitellmUserRoles.INTERNAL_USER, False), + (LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False), + ], + ) + def test_user_has_admin_view_by_role(self, user_role, expected): + """Parametrized test: admin roles return True, non-admin return False.""" + mock_auth = MagicMock() + mock_auth.user_role = user_role + assert _user_has_admin_view(mock_auth) == expected + + def test_user_has_admin_view_with_user_api_key_auth(self): + """Test with actual UserAPIKeyAuth object.""" + auth_admin = UserAPIKeyAuth( + user_id="u1", + api_key="sk-xxx", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + auth_user = UserAPIKeyAuth( + user_id="u2", + api_key="sk-yyy", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + assert _user_has_admin_view(auth_admin) is True + assert _user_has_admin_view(auth_user) is False + + +class TestIsUserTeamAdmin: + """Tests for _is_user_team_admin function.""" + + @pytest.mark.parametrize( + "members_with_roles,user_id,expected", + [ + ( + [Member(user_id="u1", role="admin")], + "u1", + True, + ), + ( + [Member(user_id="u1", role="user")], + "u1", + False, + ), + ( + [Member(user_id="u2", role="admin"), Member(user_id="u1", role="admin")], + "u1", + True, + ), + ([], "u1", False), + ], + ) + def test_is_user_team_admin_parametrized( + self, members_with_roles, user_id, expected + ): + """Parametrized test: user is team admin only when in members_with_roles with admin role.""" + mock_auth = MagicMock() + mock_auth.user_id = user_id + team = LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=members_with_roles, + ) + assert _is_user_team_admin(mock_auth, team) == expected + + def test_is_user_team_admin_user_not_in_team(self): + """Test returns False when user is not in team members.""" + auth = UserAPIKeyAuth(user_id="u99", api_key="sk-x", user_role=None) + team = LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[Member(user_id="u1", role="admin")], + ) + assert _is_user_team_admin(auth, team) is False + + +class TestOrgAdminCanInviteUser: + """Tests for _org_admin_can_invite_user function.""" + + def _make_membership(self, org_id: str, user_role: str): + now = datetime.now(timezone.utc) + return LiteLLM_OrganizationMembershipTable( + user_id="u", + organization_id=org_id, + user_role=user_role, + created_at=now, + updated_at=now, + ) + + @pytest.mark.parametrize( + "admin_orgs,target_orgs,expected", + [ + (["org1"], ["org1"], True), + (["org1", "org2"], ["org2"], True), + (["org1"], ["org2"], False), + ([], ["org1"], False), + (["org1"], [], False), + ], + ) + def test_org_admin_can_invite_user_parametrized( + self, admin_orgs, target_orgs, expected + ): + """Parametrized test: can invite when target is in org where admin has ORG_ADMIN role.""" + admin_user = LiteLLM_UserTable( + user_id="admin", + organization_memberships=[ + self._make_membership(oid, LitellmUserRoles.ORG_ADMIN.value) + for oid in admin_orgs + ], + ) + target_user = LiteLLM_UserTable( + user_id="target", + organization_memberships=[ + self._make_membership(oid, LitellmUserRoles.INTERNAL_USER.value) + for oid in target_orgs + ], + ) + assert _org_admin_can_invite_user(admin_user, target_user) == expected + + def test_org_admin_can_invite_user_no_shared_org(self): + """Test returns False when admin has no org admin role.""" + admin_user = LiteLLM_UserTable( + user_id="admin", + organization_memberships=[ + self._make_membership("org1", LitellmUserRoles.INTERNAL_USER.value), + ], + ) + target_user = LiteLLM_UserTable( + user_id="target", + organization_memberships=[ + self._make_membership("org1", LitellmUserRoles.INTERNAL_USER.value), + ], + ) + assert _org_admin_can_invite_user(admin_user, target_user) is False + + +class TestTeamAdminCanInviteUser: + """Tests for _team_admin_can_invite_user async function.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "admin_teams,target_teams,user_is_admin_in,expected", + [ + (["t1"], ["t1"], ["t1"], True), + (["t1", "t2"], ["t2"], ["t1", "t2"], True), + (["t1"], ["t2"], ["t1"], False), + ], + ) + async def test_team_admin_can_invite_user_parametrized( + self, admin_teams, target_teams, user_is_admin_in, expected + ): + """Parametrized test: can invite when target shares a team where user is admin.""" + mock_prisma = MagicMock() + mock_auth = MagicMock() + mock_auth.user_id = "admin" + + admin_user = LiteLLM_UserTable(user_id="admin", teams=admin_teams) + target_user = LiteLLM_UserTable(user_id="target", teams=target_teams) + + def make_team(tid, is_admin): + m = ( + [{"user_id": "admin", "role": "admin"}] + if is_admin + else [] + ) + obj = MagicMock() + obj.team_id = tid + obj.model_dump = lambda: {"team_id": tid, "members_with_roles": m} + return obj + + teams = [ + make_team(tid, tid in user_is_admin_in) for tid in admin_teams + ] + mock_prisma.db.litellm_teamtable.find_many = AsyncMock( + return_value=teams + ) + + result = await _team_admin_can_invite_user( + user_api_key_dict=mock_auth, + admin_user_obj=admin_user, + target_user_obj=target_user, + prisma_client=mock_prisma, + ) + assert result == expected + + @pytest.mark.asyncio + async def test_team_admin_can_invite_user_no_shared_team(self): + """Test returns False when admin and target share no team.""" + mock_prisma = MagicMock() + mock_auth = MagicMock() + mock_auth.user_id = "admin" + admin_user = LiteLLM_UserTable(user_id="admin", teams=[]) + target_user = LiteLLM_UserTable(user_id="target", teams=["t1"]) + + result = await _team_admin_can_invite_user( + user_api_key_dict=mock_auth, + admin_user_obj=admin_user, + target_user_obj=target_user, + prisma_client=mock_prisma, + ) + assert result is False + + +class TestUserHasAdminPrivileges: + """Tests for _user_has_admin_privileges async function.""" + + @pytest.mark.asyncio + async def test_proxy_admin_has_privileges(self): + """Proxy admin always has admin privileges.""" + auth = UserAPIKeyAuth( + user_id="admin", + api_key="sk-x", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + result = await _user_has_admin_privileges( + user_api_key_dict=auth, + prisma_client=None, + ) + assert result is True + + @pytest.mark.asyncio + async def test_non_admin_no_prisma_returns_false(self): + """Non-admin with no prisma connection has no privileges.""" + auth = UserAPIKeyAuth( + user_id="user1", + api_key="sk-x", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + result = await _user_has_admin_privileges( + user_api_key_dict=auth, + prisma_client=None, + ) + assert result is False + + +class TestAdminCanInviteUser: + """Tests for admin_can_invite_user async function.""" + + @pytest.mark.asyncio + async def test_proxy_admin_can_invite_any_user(self): + """Proxy admin can invite any user regardless of org/team.""" + auth = UserAPIKeyAuth( + user_id="admin", + api_key="sk-x", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + result = await admin_can_invite_user( + target_user_id="any-user", + user_api_key_dict=auth, + prisma_client=None, + ) + assert result is True + + @pytest.mark.asyncio + async def test_non_admin_cannot_invite_without_prisma(self): + """Non-admin with no prisma cannot invite.""" + auth = UserAPIKeyAuth( + user_id="user1", + api_key="sk-x", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + result = await admin_can_invite_user( + target_user_id="other-user", + user_api_key_dict=auth, + prisma_client=None, + ) + assert result is False + + +class TestSetObjectMetadataField: + """Tests for _set_object_metadata_field function.""" + + @pytest.mark.parametrize( + "field_name,value,should_call_premium", + [ + ("guardrails", ["g1"], True), + ("model_rpm_limit", {"gpt-4": 10}, False), + ], + ) + def test_set_object_metadata_field_parametrized( + self, field_name, value, should_call_premium + ): + """Parametrized test: premium fields trigger _premium_user_check.""" + team = LiteLLM_TeamTable(team_id="t1", metadata={}) + with patch( + "litellm.proxy.management_endpoints.common_utils._premium_user_check" + ) as mock_premium: + _set_object_metadata_field(team, field_name, value) + if should_call_premium: + mock_premium.assert_called_once() + else: + mock_premium.assert_not_called() + assert team.metadata[field_name] == value + + def test_set_object_metadata_field_initializes_metadata_if_none(self): + """Test initializes metadata dict when object has None.""" + team = LiteLLM_TeamTable(team_id="t1", metadata=None) + with patch( + "litellm.proxy.management_endpoints.common_utils._premium_user_check" + ): + _set_object_metadata_field(team, "model_rpm_limit", {"x": 1}) + assert team.metadata == {"model_rpm_limit": {"x": 1}} diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 6331e99462b..e81c6264f7b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1,4 +1,3 @@ -import json import os import sys import types @@ -210,18 +209,23 @@ class TestListMCPServers: for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", - AsyncMock(return_value=[mock_user_auth]), + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -268,12 +272,15 @@ class TestListMCPServers: return_value=mock_servers ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", - return_value="view_all", - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", + return_value="view_all", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), ): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( fetch_all_mcp_servers, @@ -284,6 +291,79 @@ class TestListMCPServers: assert len(result) == 2 assert {server.server_id for server in result} == {"server-1", "server-2"} + @pytest.mark.asyncio + async def test_list_mcp_servers_view_all_mode_virtual_key_is_sanitized(self): + """Issue #20325: virtual keys should get a safe discovery view.""" + + mock_user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test_user_id", + api_key="test_api_key", + allowed_routes=["mcp_routes"], + ) + + mock_servers = [ + generate_mock_mcp_server_db_record(server_id="server-1", alias="One"), + generate_mock_mcp_server_db_record(server_id="server-2", alias="Two"), + ] + for idx, server in enumerate(mock_servers): + server.credentials = {"auth_value": f"secret_{idx}"} + server.env = {"API_KEY": "super-secret"} + server.static_headers = {"Authorization": "Bearer super-secret"} + server.mcp_access_groups = ["group-a"] + server.teams = [{"team_id": "team-1", "team_alias": "Team 1"}] + server.command = "bash" + server.args = ["-lc", "echo hi"] + server.extra_headers = ["Authorization"] + + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_unfiltered = AsyncMock( + return_value=mock_servers + ) + mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=mock_servers) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", + return_value="view_all", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth) + + # Ensure we did not bypass filtering via view_all for restricted virtual keys. + mock_manager.get_all_mcp_servers_unfiltered.assert_not_called() + + assert len(result) == 2 + assert {server.server_id for server in result} == {"server-1", "server-2"} + + for server in result: + assert server.credentials is None + assert server.url is None + assert server.static_headers is None + assert server.env == {} + assert server.command is None + assert server.args == [] + assert server.extra_headers == [] + assert server.allowed_tools == [] + assert server.mcp_access_groups == [] + assert server.teams == [] + @pytest.mark.asyncio async def test_list_mcp_servers_combined_config_and_db(self): """ @@ -378,18 +458,23 @@ class TestListMCPServers: for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", - AsyncMock(return_value=[mock_user_auth]), + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -496,18 +581,23 @@ class TestListMCPServers: for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", - AsyncMock(return_value=[mock_user_auth]), + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -622,18 +712,23 @@ class TestListMCPServers: user_role=LitellmUserRoles.PROXY_ADMIN ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=mock_server), - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", - AsyncMock(return_value=mock_health_result), - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=mock_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", + AsyncMock(return_value=mock_health_result), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), ): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( fetch_mcp_server, @@ -670,18 +765,23 @@ class TestListMCPServers: user_role=LitellmUserRoles.PROXY_ADMIN ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=mock_server), - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", - AsyncMock(return_value=mock_health_result), - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=mock_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", + AsyncMock(return_value=mock_health_result), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), ): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( fetch_mcp_server, @@ -823,16 +923,20 @@ class TestTemporaryMCPSessionEndpoints: mock_manager.get_mcp_server_by_id.return_value = inherited_server mock_manager.build_mcp_server_from_table = AsyncMock(return_value=built_server) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", - MagicMock(), - ) as validate_mock, patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._cache_temporary_mcp_server", - MagicMock(), - ) as cache_mock: + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + MagicMock(), + ) as validate_mock, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._cache_temporary_mcp_server", + MagicMock(), + ) as cache_mock, + ): response = await add_session_mcp_server( payload=payload, user_api_key_dict=user_auth, @@ -892,13 +996,16 @@ class TestTemporaryMCPSessionEndpoints: server = generate_mock_mcp_server_config_record(server_id="server-1") authorize_response = MagicMock() - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, - ) as get_server, patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.authorize_with_server", - AsyncMock(return_value=authorize_response), - ) as authorize_mock: + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ) as get_server, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.authorize_with_server", + AsyncMock(return_value=authorize_response), + ) as authorize_mock, + ): result = await mcp_authorize( request=request, server_id="server-1", @@ -935,13 +1042,16 @@ class TestTemporaryMCPSessionEndpoints: server = generate_mock_mcp_server_config_record(server_id="server-1") exchange_response = {"access_token": "token"} - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, - ) as get_server, patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", - AsyncMock(return_value=exchange_response), - ) as exchange_mock: + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ) as get_server, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", + AsyncMock(return_value=exchange_response), + ) as exchange_mock, + ): result = await mcp_token( request=request, server_id="server-1", @@ -982,16 +1092,20 @@ class TestTemporaryMCPSessionEndpoints: "token_endpoint_auth_method": "client_secret_basic", } - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, - ) as get_server, patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", - AsyncMock(return_value=request_body), - ) as read_body, patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.register_client_with_server", - AsyncMock(return_value=register_response), - ) as register_mock: + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ) as get_server, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", + AsyncMock(return_value=request_body), + ) as read_body, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.register_client_with_server", + AsyncMock(return_value=register_response), + ) as register_mock, + ): result = await mcp_register(request=request, server_id="server-1") assert result is register_response @@ -1060,21 +1174,27 @@ class TestUpdateMCPServer: ) # Mock the update_mcp_server function to capture the call - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", - MagicMock(), - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", - AsyncMock(return_value=updated_server), - ) as update_mock, patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_server", - AsyncMock(), - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.reload_servers_from_database", - AsyncMock(), + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", + AsyncMock(return_value=updated_server), + ) as update_mock, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_server", + AsyncMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.reload_servers_from_database", + AsyncMock(), + ), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -1145,12 +1265,15 @@ class TestHealthCheckServers: return_value=[mock_health_result_1, mock_health_result_2] ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", - AsyncMock(return_value=[mock_user_auth]), + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), ): result = await health_check_servers( server_ids=None, @@ -1199,9 +1322,12 @@ class TestMCPRegistryEndpoint: mock_server.server_id: mock_server } - with patch_proxy_general_settings({"enable_mcp_registry": True}), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, + with ( + patch_proxy_general_settings({"enable_mcp_registry": True}), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), ): response = client.get("/v1/mcp/registry.json") @@ -1248,12 +1374,15 @@ class TestMCPRegistryEndpoint: return_value=[mock_health_result] ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", - AsyncMock(return_value=[mock_user_auth]), + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), ): result = await health_check_servers( server_ids=["server-1"], @@ -1311,12 +1440,15 @@ class TestManagementPayloadValidation: return_value=[health_result_one, health_result_two] ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", - return_value="view_all", - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", + return_value="view_all", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), ): result = await health_check_servers( server_ids=None, @@ -1361,12 +1493,15 @@ class TestManagementPayloadValidation: return_value=[mock_health_result] # Only server-1 is returned (accessible) ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", - AsyncMock(return_value=[mock_user_auth]), + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), ): result = await health_check_servers( server_ids=["server-1", "server-unauthorized"], diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 16f80826798..74d36c0acac 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2374,47 +2374,6 @@ class TestProcessSSOJWTAccessToken: "groups": ["team1", "team2", "team3"], } - def test_process_sso_jwt_access_token_with_valid_token( - self, mock_jwt_handler, sample_jwt_token, sample_jwt_payload - ): - """Test processing a valid JWT access token with team extraction""" - from litellm.proxy.management_endpoints.ui_sso import ( - process_sso_jwt_access_token, - ) - - # Create a result object without team_ids - result = CustomOpenID( - id="test_user", - email="test@example.com", - first_name="Test", - last_name="User", - display_name="Test User", - provider="generic", - team_ids=[], - ) - - with patch("jwt.decode", return_value=sample_jwt_payload) as mock_jwt_decode: - # Act - process_sso_jwt_access_token( - access_token_str=sample_jwt_token, - sso_jwt_handler=mock_jwt_handler, - result=result, - ) - - # Assert - # Verify JWT was decoded correctly - mock_jwt_decode.assert_called_once_with( - sample_jwt_token, options={"verify_signature": False} - ) - - # Verify team IDs were extracted from JWT - mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with( - sample_jwt_payload - ) - - # Verify team IDs were set on the result object - assert result.team_ids == ["team1", "team2", "team3"] - def test_process_sso_jwt_access_token_with_existing_team_ids( self, mock_jwt_handler, sample_jwt_token ): @@ -2549,27 +2508,6 @@ class TestProcessSSOJWTAccessToken: mock_jwt_handler.get_team_ids_from_jwt.assert_not_called() assert result.team_ids == [] - def test_process_sso_jwt_access_token_no_sso_jwt_handler(self, sample_jwt_token): - """Test that JWT is decoded for role extraction even when sso_jwt_handler is None, - but team_ids are not extracted (team extraction requires sso_jwt_handler).""" - from litellm.proxy.management_endpoints.ui_sso import ( - process_sso_jwt_access_token, - ) - - result = CustomOpenID(id="test_user", email="test@example.com", team_ids=[]) - - mock_payload = {"sub": "test_user", "email": "test@example.com"} - with patch("jwt.decode", return_value=mock_payload) as mock_jwt_decode: - # Act - process_sso_jwt_access_token( - access_token_str=sample_jwt_token, sso_jwt_handler=None, result=result - ) - - # JWT is decoded (for role extraction) but team_ids are not extracted - mock_jwt_decode.assert_called_once() - assert result.team_ids == [] - assert result.user_role is None - def test_process_sso_jwt_access_token_no_result( self, mock_jwt_handler, sample_jwt_token ): @@ -2590,10 +2528,12 @@ class TestProcessSSOJWTAccessToken: mock_jwt_decode.assert_not_called() mock_jwt_handler.get_team_ids_from_jwt.assert_not_called() - def test_process_sso_jwt_access_token_jwt_decode_exception( + def test_process_sso_jwt_access_token_non_decode_exception_propagates( self, mock_jwt_handler, sample_jwt_token ): - """Test that JWT decode exceptions are not caught (should propagate up)""" + """Test that non-DecodeError JWT exceptions still propagate up.""" + import jwt as pyjwt + from litellm.proxy.management_endpoints.ui_sso import ( process_sso_jwt_access_token, ) @@ -2601,19 +2541,16 @@ class TestProcessSSOJWTAccessToken: result = CustomOpenID(id="test_user", email="test@example.com", team_ids=[]) with patch( - "jwt.decode", side_effect=Exception("JWT decode error") + "jwt.decode", side_effect=pyjwt.exceptions.InvalidKeyError("Invalid key") ) as mock_jwt_decode: - # Act & Assert - with pytest.raises(Exception, match="JWT decode error"): + with pytest.raises(pyjwt.exceptions.InvalidKeyError, match="Invalid key"): process_sso_jwt_access_token( access_token_str=sample_jwt_token, sso_jwt_handler=mock_jwt_handler, result=result, ) - # Verify JWT decode was attempted mock_jwt_decode.assert_called_once() - # But team extraction should not have been called mock_jwt_handler.get_team_ids_from_jwt.assert_not_called() def test_process_sso_jwt_access_token_empty_team_ids_from_jwt( @@ -2646,6 +2583,124 @@ class TestProcessSSOJWTAccessToken: # Even empty team IDs should be set assert result.team_ids == [] + def test_process_sso_jwt_access_token_with_opaque_token(self, mock_jwt_handler): + """Test that opaque (non-JWT) access tokens are handled gracefully without raising.""" + from litellm.proxy.management_endpoints.ui_sso import ( + process_sso_jwt_access_token, + ) + + result = CustomOpenID( + id="test_user", + email="test@example.com", + first_name="Test", + last_name="User", + display_name="Test User", + provider="generic", + team_ids=["existing_team"], + user_role=None, + ) + + # Opaque tokens like those from Logto are short random strings, not JWTs + opaque_token = "uTxyjXbS_random_opaque_token_string" + + # Should NOT raise - opaque tokens should be silently skipped + process_sso_jwt_access_token( + access_token_str=opaque_token, + sso_jwt_handler=mock_jwt_handler, + result=result, + ) + + # Result should be untouched + mock_jwt_handler.get_team_ids_from_jwt.assert_not_called() + assert result.team_ids == ["existing_team"] + assert result.user_role is None + + def test_process_sso_jwt_access_token_real_jwt_with_role_and_teams( + self, mock_jwt_handler + ): + """Test that a real JWT containing role and team fields is correctly processed.""" + import jwt as pyjwt + + from litellm.proxy.management_endpoints.ui_sso import ( + process_sso_jwt_access_token, + ) + + payload = { + "sub": "user123", + "email": "admin@example.com", + "role": "proxy_admin", + "groups": ["team_alpha", "team_beta"], + } + real_jwt_token = pyjwt.encode(payload, "test-secret", algorithm="HS256") + + mock_jwt_handler.get_team_ids_from_jwt.return_value = [ + "team_alpha", + "team_beta", + ] + + result = CustomOpenID( + id="user123", + email="admin@example.com", + first_name="Admin", + last_name="User", + display_name="Admin User", + provider="generic", + team_ids=[], + user_role=None, + ) + + process_sso_jwt_access_token( + access_token_str=real_jwt_token, + sso_jwt_handler=mock_jwt_handler, + result=result, + ) + + # Team IDs should be extracted via sso_jwt_handler + mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with(payload) + assert result.team_ids == ["team_alpha", "team_beta"] + + # Role should be extracted from the "role" field in the JWT + from litellm.proxy._types import LitellmUserRoles + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + def test_process_sso_jwt_access_token_real_jwt_without_role_and_teams(self): + """Test that a real JWT without role/team fields leaves result unchanged.""" + import jwt as pyjwt + + from litellm.proxy.management_endpoints.ui_sso import ( + process_sso_jwt_access_token, + ) + + payload = { + "sub": "user456", + "email": "plain@example.com", + "iat": 1700000000, + } + real_jwt_token = pyjwt.encode(payload, "test-secret", algorithm="HS256") + + result = CustomOpenID( + id="user456", + email="plain@example.com", + first_name="Plain", + last_name="User", + display_name="Plain User", + provider="generic", + team_ids=[], + user_role=None, + ) + + # No sso_jwt_handler, no role/team fields in JWT + process_sso_jwt_access_token( + access_token_str=real_jwt_token, + sso_jwt_handler=None, + result=result, + ) + + # Nothing should be modified + assert result.team_ids == [] + assert result.user_role is None + @pytest.mark.asyncio async def test_get_ui_settings_includes_api_doc_base_url(): @@ -4071,3 +4126,123 @@ def test_process_sso_jwt_access_token_with_role_mappings(): # Should get highest privilege role assert result.user_role == LitellmUserRoles.PROXY_ADMIN + +def test_generic_response_convertor_with_extra_attributes(monkeypatch): + """Test that extra attributes are extracted when GENERIC_USER_EXTRA_ATTRIBUTES is set""" + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + monkeypatch.setenv("GENERIC_CLIENT_ID", "test_client") + monkeypatch.setenv("GENERIC_USER_EXTRA_ATTRIBUTES", "custom_field1,custom_field2,custom_field3") + + mock_response = { + "sub": "user-id-123", + "email": "user@example.com", + "given_name": "John", + "family_name": "Doe", + "name": "John Doe", + "provider": "generic", + "custom_field1": "value1", + "custom_field2": ["item1", "item2"], + "custom_field3": {"nested": "data"}, + } + + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + result = generic_response_convertor( + response=mock_response, + jwt_handler=mock_jwt_handler, + sso_jwt_handler=None, + role_mappings=None, + ) + + assert result.extra_fields is not None + assert result.extra_fields["custom_field1"] == "value1" + assert result.extra_fields["custom_field2"] == ["item1", "item2"] + assert result.extra_fields["custom_field3"] == {"nested": "data"} + +def test_generic_response_convertor_without_extra_attributes(monkeypatch): + """Test backward compatibility - extra_fields is None when env var not set""" + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + monkeypatch.setenv("GENERIC_CLIENT_ID", "test_client") + # Don't set GENERIC_USER_EXTRA_ATTRIBUTES + + mock_response = { + "sub": "user-id-123", + "email": "user@example.com", + "given_name": "John", + "family_name": "Doe", + "name": "John Doe", + "provider": "generic", + "custom_field1": "value1", + "custom_field2": "value2", + } + + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + result = generic_response_convertor( + response=mock_response, + jwt_handler=mock_jwt_handler, + sso_jwt_handler=None, + role_mappings=None, + ) + + assert result.extra_fields is None + +def test_generic_response_convertor_extra_attributes_with_nested_paths(monkeypatch): + """Test that nested paths work with dot notation""" + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + monkeypatch.setenv("GENERIC_CLIENT_ID", "test_client") + monkeypatch.setenv("GENERIC_USER_EXTRA_ATTRIBUTES", "org_info.department,org_info.manager") + + mock_response = { + "sub": "user-id-123", + "email": "user@example.com", + "org_info": { + "department": "Engineering", + "manager": "Jane Smith" + } + } + + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + result = generic_response_convertor( + response=mock_response, + jwt_handler=mock_jwt_handler, + sso_jwt_handler=None, + role_mappings=None, + ) + + assert result.extra_fields is not None + assert result.extra_fields["org_info.department"] == "Engineering" + assert result.extra_fields["org_info.manager"] == "Jane Smith" + +def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch): + """Test that missing fields return None""" + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + monkeypatch.setenv("GENERIC_CLIENT_ID", "test_client") + monkeypatch.setenv("GENERIC_USER_EXTRA_ATTRIBUTES", "missing_field,another_missing") + + mock_response = { + "sub": "user-id-123", + "email": "user@example.com", + } + + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + result = generic_response_convertor( + response=mock_response, + jwt_handler=mock_jwt_handler, + sso_jwt_handler=None, + role_mappings=None, + ) + + assert result.extra_fields is not None + assert result.extra_fields["missing_field"] is None + assert result.extra_fields["another_missing"] is None \ No newline at end of file diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index daae6d465a7..e50e10352e2 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1316,6 +1316,133 @@ async def test_delete_pass_through_endpoint_not_found(): assert "not found" in str(exc_info.value.detail).lower() +@pytest.mark.asyncio +async def test_get_pass_through_endpoints_includes_config_and_db(): + """ + Test that get_pass_through_endpoints returns both config-defined and DB endpoints, + with correct is_from_config flag. Config-only endpoints have is_from_config=True, + DB endpoints have is_from_config=False. When same path exists in both, DB overrides. + """ + from litellm.proxy._types import ( + PassThroughEndpointResponse, + PassThroughGenericEndpoint, + UserAPIKeyAuth, + ) + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + get_pass_through_endpoints, + ) + + # Config-defined endpoints (from config file) + config_endpoints = [ + { + "path": "/v1/rerank", + "target": "https://api.cohere.com/v1/rerank", + "headers": {"content-type": "application/json"}, + }, + { + "path": "/v1/config-only", + "target": "https://config.example.com/api", + "headers": {}, + }, + ] + + # DB endpoints (one overlaps with config path, one is DB-only) + db_endpoints = [ + { + "id": "db-endpoint-1", + "path": "/v1/rerank", # Same as config - DB should override + "target": "https://db-override.com/v1/rerank", + "headers": {}, + "include_subpath": False, + }, + { + "id": "db-endpoint-2", + "path": "/db/only", + "target": "https://db-only.example.com/api", + "headers": {}, + "include_subpath": False, + }, + ] + + with patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ): + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._get_pass_through_endpoints_from_db", + new_callable=AsyncMock, + ) as mock_get_db: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._get_pass_through_endpoints_from_config" + ) as mock_get_config: + db_objects = [ + PassThroughGenericEndpoint(**ep, is_from_config=False) + for ep in db_endpoints + ] + config_objects = [ + PassThroughGenericEndpoint(**ep, is_from_config=True) + for ep in config_endpoints + ] + mock_get_db.return_value = db_objects + mock_get_config.return_value = config_objects + + mock_user = MagicMock(spec=UserAPIKeyAuth) + + result = await get_pass_through_endpoints( + endpoint_id=None, + user_api_key_dict=mock_user, + team_id=None, + ) + + assert isinstance(result, PassThroughEndpointResponse) + # config_only: /v1/config-only (not in db_paths) + # db: /v1/rerank (overrides config), /db/only + # So we should have: /v1/config-only (from config) + /v1/rerank + /db/only (from db) + assert len(result.endpoints) == 3 + + # Check is_from_config values + by_path = {ep.path: ep for ep in result.endpoints} + assert by_path["/v1/config-only"].is_from_config is True + assert by_path["/v1/rerank"].is_from_config is False # DB overrides + assert by_path["/db/only"].is_from_config is False + + # Verify DB override: /v1/rerank should have DB target + assert by_path["/v1/rerank"].target == "https://db-override.com/v1/rerank" + + +def test_get_pass_through_endpoints_from_config_skips_malformed(): + """ + Test that _get_pass_through_endpoints_from_config skips malformed endpoints + and returns only valid ones, without raising. + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _get_pass_through_endpoints_from_config, + ) + + # Mix of valid and malformed config endpoints + config_passthrough_endpoints = [ + {"path": "/valid/1", "target": "https://valid1.example.com"}, + {}, # Missing required path and target + {"path": "/missing-target"}, # Missing required target + {"target": "https://example.com"}, # Missing required path + {"path": "/valid/2", "target": "https://valid2.example.com", "headers": {}}, + ] + + with patch( + "litellm.proxy.proxy_server.config_passthrough_endpoints", + config_passthrough_endpoints, + ): + result = _get_pass_through_endpoints_from_config() + + # Only the 2 valid endpoints should be returned + assert len(result) == 2 + paths = {ep.path for ep in result} + assert "/valid/1" in paths + assert "/valid/2" in paths + for ep in result: + assert ep.is_from_config is True + + @pytest.mark.asyncio async def test_delete_pass_through_endpoint_empty_list(): """ diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 1ed956fe99f..c853253eedd 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -192,6 +192,139 @@ class TestGetAttachedPolicies: assert "strict-policy" not in registry.get_attached_policies(context_wrong_team) +class TestTagBasedAttachments: + """Test tag-based policy attachment matching.""" + + def test_tag_matching_and_wildcards(self): + """Test tag matching: exact match, wildcard match, and no-match cases.""" + registry = AttachmentRegistry() + registry.load_attachments([ + {"policy": "hipaa-policy", "tags": ["healthcare"]}, + {"policy": "health-policy", "tags": ["health-*"]}, + ]) + + # Exact tag match + context = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4", + tags=["healthcare"], + ) + attached = registry.get_attached_policies(context) + assert "hipaa-policy" in attached + assert "health-policy" not in attached # "healthcare" doesn't match "health-*" + + # Wildcard tag match + context_wildcard = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4", + tags=["health-prod"], + ) + attached_wildcard = registry.get_attached_policies(context_wildcard) + assert "health-policy" in attached_wildcard + assert "hipaa-policy" not in attached_wildcard + + # No match — wrong tag + context_no_match = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4", + tags=["finance"], + ) + assert registry.get_attached_policies(context_no_match) == [] + + # No match — no tags on context + context_no_tags = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4", + tags=None, + ) + assert registry.get_attached_policies(context_no_tags) == [] + + def test_tag_combined_with_team(self): + """Test attachment with both tags and teams requires BOTH to match (AND logic).""" + registry = AttachmentRegistry() + registry.load_attachments([ + {"policy": "strict-policy", "teams": ["team-a"], "tags": ["healthcare"]}, + ]) + + # Match — both team and tag match + context = PolicyMatchContext( + team_alias="team-a", key_alias="key", model="gpt-4", + tags=["healthcare"], + ) + assert "strict-policy" in registry.get_attached_policies(context) + + # No match — tag matches but team doesn't + context_wrong_team = PolicyMatchContext( + team_alias="team-b", key_alias="key", model="gpt-4", + tags=["healthcare"], + ) + assert "strict-policy" not in registry.get_attached_policies(context_wrong_team) + + # No match — team matches but tag doesn't + context_wrong_tag = PolicyMatchContext( + team_alias="team-a", key_alias="key", model="gpt-4", + tags=["finance"], + ) + assert "strict-policy" not in registry.get_attached_policies(context_wrong_tag) + + +class TestMatchAttribution: + """Test get_attached_policies_with_reasons — the attribution logic that + powers response headers and the Policy Simulator UI.""" + + def test_reasons_for_global_tag_team_attachments(self): + """Test that match reasons correctly describe WHY each policy matched.""" + registry = AttachmentRegistry() + registry.load_attachments([ + {"policy": "global-baseline", "scope": "*"}, + {"policy": "hipaa-policy", "tags": ["healthcare"]}, + {"policy": "team-policy", "teams": ["health-team"]}, + ]) + + context = PolicyMatchContext( + team_alias="health-team", key_alias="key", model="gpt-4", + tags=["healthcare"], + ) + results = registry.get_attached_policies_with_reasons(context) + reasons = {r["policy_name"]: r["matched_via"] for r in results} + + assert reasons["global-baseline"] == "scope:*" + assert "tag:healthcare" in reasons["hipaa-policy"] + assert "team:health-team" in reasons["team-policy"] + + def test_tags_only_attachment_matches_any_team_key_model(self): + """Test the primary use case: tags-only attachment with no team/key/model + constraint matches any request that carries the tag.""" + registry = AttachmentRegistry() + registry.load_attachments([ + {"policy": "hipaa-guardrails", "tags": ["healthcare"]}, + ]) + + # Should match regardless of team/key/model + context = PolicyMatchContext( + team_alias="random-team", key_alias="random-key", model="claude-3", + tags=["healthcare"], + ) + attached = registry.get_attached_policies(context) + assert "hipaa-guardrails" in attached + + # Should not match without the tag + context_no_tag = PolicyMatchContext( + team_alias="random-team", key_alias="random-key", model="claude-3", + ) + assert registry.get_attached_policies(context_no_tag) == [] + + def test_attachment_with_no_scope_matches_everything(self): + """Test that an attachment with no scope/teams/keys/models/tags + matches everything because teams/keys/models default to ['*'].""" + registry = AttachmentRegistry() + registry.load_attachments([ + {"policy": "catch-all"}, + ]) + + context = PolicyMatchContext( + team_alias="any-team", key_alias="any-key", model="gpt-4", + ) + attached = registry.get_attached_policies(context) + assert "catch-all" in attached + + class TestAttachmentRegistrySingleton: """Test global singleton behavior.""" diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py b/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py index c011f31af6a..fccb26496ac 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py @@ -64,6 +64,70 @@ class TestPolicyMatcherScopeMatching: assert PolicyMatcher.scope_matches(scope, context) is True +class TestPolicyMatcherScopeMatchingWithTags: + """Test scope matching with tag patterns.""" + + def test_scope_tag_matching(self): + """Test scope tag matching: exact, wildcard, no-match, and empty context tags.""" + # Exact match + scope = PolicyScope(teams=["*"], keys=["*"], models=["*"], tags=["healthcare"]) + context = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4", + tags=["healthcare", "internal"], + ) + assert PolicyMatcher.scope_matches(scope, context) is True + + # Wildcard match + scope_wc = PolicyScope(teams=["*"], keys=["*"], models=["*"], tags=["health-*"]) + context_wc = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4", + tags=["health-prod"], + ) + assert PolicyMatcher.scope_matches(scope_wc, context_wc) is True + + # No match — wrong tag + context_wrong = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4", + tags=["finance"], + ) + assert PolicyMatcher.scope_matches(scope, context_wrong) is False + + # No match — context has no tags + context_none = PolicyMatchContext( + team_alias="team", key_alias="key", model="gpt-4", tags=None, + ) + assert PolicyMatcher.scope_matches(scope, context_none) is False + + # Scope without tags matches any context (opt-in semantics) + scope_no_tags = PolicyScope(teams=["*"], keys=["*"], models=["*"]) + assert PolicyMatcher.scope_matches(scope_no_tags, context) is True + + def test_scope_tags_and_team_combined(self): + """Test scope with both tags and team — both must match (AND logic).""" + scope = PolicyScope(teams=["team-a"], keys=["*"], models=["*"], tags=["healthcare"]) + + # Both match + context_both = PolicyMatchContext( + team_alias="team-a", key_alias="key", model="gpt-4", + tags=["healthcare"], + ) + assert PolicyMatcher.scope_matches(scope, context_both) is True + + # Tag matches, team doesn't + context_wrong_team = PolicyMatchContext( + team_alias="team-b", key_alias="key", model="gpt-4", + tags=["healthcare"], + ) + assert PolicyMatcher.scope_matches(scope, context_wrong_team) is False + + # Team matches, tag doesn't + context_wrong_tag = PolicyMatchContext( + team_alias="team-a", key_alias="key", model="gpt-4", + tags=["finance"], + ) + assert PolicyMatcher.scope_matches(scope, context_wrong_tag) is False + + class TestPolicyMatcherWithAttachments: """Test getting matching policies via attachments.""" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 54a276bc97f..0d4a711812f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -12,7 +12,7 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import litellm import litellm.proxy.proxy_server as ps @@ -206,6 +206,7 @@ ignored_keys = [ "metadata.additional_usage_values.cache_creation_input_tokens", "metadata.additional_usage_values.cache_read_input_tokens", "metadata.additional_usage_values.inference_geo", + "metadata.additional_usage_values.speed", "metadata.litellm_overhead_time_ms", "metadata.cost_breakdown", ] @@ -346,6 +347,196 @@ async def test_ui_view_spend_logs_with_user_id(client, monkeypatch): assert data["data"][0]["user"] == "test_user_1" +# Mock spend logs with distinct values for sorting tests. +# req_a: spend=0.10, tokens=500, start/end earliest +# req_b: spend=0.05, tokens=200, start/end 2nd +# req_c: spend=0.20, tokens=50, start/end latest +# req_d: spend=0.01, tokens=100, start/end 3rd +_SORT_TEST_LOGS = [ + { + "request_id": "req_a", + "api_key": "sk-test-key", + "user": "user1", + "spend": 0.10, + "total_tokens": 500, + "startTime": "2025-01-01T00:00:00+00:00", + "endTime": "2025-01-01T00:01:00+00:00", + "model": "gpt-3.5-turbo", + }, + { + "request_id": "req_b", + "api_key": "sk-test-key", + "user": "user1", + "spend": 0.05, + "total_tokens": 200, + "startTime": "2025-01-01T00:00:01+00:00", + "endTime": "2025-01-01T00:01:01+00:00", + "model": "gpt-3.5-turbo", + }, + { + "request_id": "req_c", + "api_key": "sk-test-key", + "user": "user1", + "spend": 0.20, + "total_tokens": 50, + "startTime": "2025-01-01T00:00:03+00:00", + "endTime": "2025-01-01T00:01:03+00:00", + "model": "gpt-3.5-turbo", + }, + { + "request_id": "req_d", + "api_key": "sk-test-key", + "user": "user1", + "spend": 0.01, + "total_tokens": 100, + "startTime": "2025-01-01T00:00:02+00:00", + "endTime": "2025-01-01T00:01:02+00:00", + "model": "gpt-3.5-turbo", + }, +] + + +def _sort_logs(logs, order_clause): + """Sort logs by the given Prisma-style order clause, e.g. {'spend': 'asc'}.""" + if not order_clause: + return list(logs) + key, direction = next(iter(order_clause.items())) + reverse = direction.lower() == "desc" + return sorted(logs, key=lambda x: x.get(key, 0), reverse=reverse) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sort_by,sort_order,expected_request_ids", + [ + # spend: 0.01(d) < 0.05(b) < 0.10(a) < 0.20(c) + ("spend", "asc", ["req_d", "req_b", "req_a", "req_c"]), + ("spend", "desc", ["req_c", "req_a", "req_b", "req_d"]), + # total_tokens: 50(c) < 100(d) < 200(b) < 500(a) + ("total_tokens", "asc", ["req_c", "req_d", "req_b", "req_a"]), + ("total_tokens", "desc", ["req_a", "req_b", "req_d", "req_c"]), + # startTime: 00:00:00(a) < 00:00:01(b) < 00:00:02(d) < 00:00:03(c) + ("startTime", "asc", ["req_a", "req_b", "req_d", "req_c"]), + ("startTime", "desc", ["req_c", "req_d", "req_b", "req_a"]), + # endTime: same ordering as startTime + ("endTime", "asc", ["req_a", "req_b", "req_d", "req_c"]), + ("endTime", "desc", ["req_c", "req_d", "req_b", "req_a"]), + # default when sort_by not provided: startTime desc + (None, "desc", ["req_c", "req_d", "req_b", "req_a"]), + ], +) +async def test_ui_view_spend_logs_sort_by_and_sort_order( + client, monkeypatch, sort_by, sort_order, expected_request_ids +): + """Test that spend logs are returned in the correct order for each sort_by/sort_order.""" + base_logs = list(_SORT_TEST_LOGS) + + async def mock_find_many(*args, **kwargs): + order = kwargs.get("order", {}) + return _sort_logs(base_logs, order) + + async def mock_count(*args, **kwargs): + return len(base_logs) + + class MockPrismaClient: + def __init__(self): + self.db = MagicMock() + self.db.litellm_spendlogs = MagicMock() + self.db.litellm_spendlogs.find_many = AsyncMock(side_effect=mock_find_many) + self.db.litellm_spendlogs.count = AsyncMock(side_effect=mock_count) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient()) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + start_date = "2024-12-25 00:00:00" + end_date = "2025-01-02 23:59:59" + + params = { + "start_date": start_date, + "end_date": end_date, + } + if sort_by is not None: + params["sort_by"] = sort_by + if sort_order is not None: + params["sort_order"] = sort_order + + response = client.get( + "/spend/logs/ui", + params=params, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200, response.text + data = response.json() + assert "data" in data + + actual_ids = [log["request_id"] for log in data["data"]] + assert actual_ids == expected_request_ids, ( + f"Expected order {expected_request_ids}, got {actual_ids} " + f"(sort_by={sort_by}, sort_order={sort_order})" + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sort_by,sort_order", + [ + ("invalid", "asc"), + ("spend", "invalid"), + ], +) +async def test_ui_view_spend_logs_sort_validation_errors( + client, monkeypatch, sort_by, sort_order +): + """Test that invalid sort_by and sort_order return 400.""" + async def mock_count(*args, **kwargs): + return 0 + + class MockPrismaClient: + def __init__(self): + self.db = MagicMock() + self.db.litellm_spendlogs = MagicMock() + self.db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) + self.db.litellm_spendlogs.count = AsyncMock(side_effect=mock_count) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient()) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + start_date = "2024-12-25 00:00:00" + end_date = "2025-01-02 23:59:59" + + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "sort_by": sort_by, + "sort_order": sort_order, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 400 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_with_team_id(client, monkeypatch): # Mock data for the test @@ -1025,6 +1216,85 @@ async def test_ui_view_spend_logs_with_model(client, monkeypatch): assert data["data"][0]["model"] == "gpt-3.5-turbo" +@pytest.mark.asyncio +async def test_ui_view_spend_logs_with_model_id(client, monkeypatch): + """Test that the model_id query param filters spend logs by litellm model deployment id.""" + mock_spend_logs = [ + { + "id": "log1", + "request_id": "req1", + "api_key": "sk-test-key", + "user": "test_user_1", + "team_id": "team1", + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-3.5-turbo", + "model_id": "deployment-id-1", + "status": "success", + }, + { + "id": "log2", + "request_id": "req2", + "api_key": "sk-test-key", + "user": "test_user_2", + "team_id": "team1", + "spend": 0.10, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + "model_id": "deployment-id-2", + "status": "success", + }, + ] + + class MockDB: + async def find_many(self, *args, **kwargs): + if ( + "where" in kwargs + and "model_id" in kwargs["where"] + and kwargs["where"]["model_id"] == "deployment-id-1" + ): + return [mock_spend_logs[0]] + return mock_spend_logs + + async def count(self, *args, **kwargs): + if ( + "where" in kwargs + and "model_id" in kwargs["where"] + and kwargs["where"]["model_id"] == "deployment-id-1" + ): + return 1 + return len(mock_spend_logs) + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + self.db.litellm_spendlogs = self.db + + mock_prisma_client = MockPrismaClient() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + start_date = ( + datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) + ).strftime("%Y-%m-%d %H:%M:%S") + end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + + response = client.get( + "/spend/logs/ui", + params={ + "model_id": "deployment-id-1", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["model_id"] == "deployment-id-1" + + @pytest.mark.asyncio async def test_ui_view_spend_logs_with_key_hash(client, monkeypatch): # Mock data for the test diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index acd99090397..e7da4256182 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3203,3 +3203,123 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch): assert result["general_settings"]["nested"]["key1"] == "updated_value1" assert result["general_settings"]["nested"]["key2"] == "value2" assert result["general_settings"]["nested"]["key3"] == "value3" + + +class TestInvitationEndpoints: + """Tests for /invitation/new and /invitation/delete endpoints.""" + + @pytest.fixture + def client_with_auth(self): + """Create a test client with admin authentication.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.proxy_server import cleanup_router_config_variables + + cleanup_router_config_variables() + filepath = os.path.dirname(os.path.abspath(__file__)) + config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" + asyncio.run(initialize(config=config_fp, debug=True)) + + mock_auth = MagicMock() + mock_auth.user_id = "admin-user-id" + mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN + mock_auth.api_key = "sk-test" + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + + return TestClient(app) + + @pytest.mark.parametrize( + "endpoint,payload,mock_return", + [ + ( + "/invitation/new", + {"user_id": "target-user-123"}, + { + "id": "inv-123", + "user_id": "target-user-123", + "is_accepted": False, + "accepted_at": None, + "expires_at": "2025-02-18T00:00:00", + "created_at": "2025-02-11T00:00:00", + "created_by": "admin-user-id", + "updated_at": "2025-02-11T00:00:00", + "updated_by": "admin-user-id", + }, + ), + ( + "/invitation/delete", + {"invitation_id": "inv-456"}, + { + "id": "inv-456", + "user_id": "target-user-123", + "is_accepted": False, + "accepted_at": None, + "expires_at": "2025-02-18T00:00:00", + "created_at": "2025-02-11T00:00:00", + "created_by": "admin-user-id", + "updated_at": "2025-02-11T00:00:00", + "updated_by": "admin-user-id", + }, + ), + ], + ) + def test_invitation_endpoints_proxy_admin_success( + self, client_with_auth, endpoint, payload, mock_return + ): + """Proxy admin can successfully create and delete invitations.""" + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_invitationlink = MagicMock() + if endpoint == "/invitation/new": + mock_create = AsyncMock(return_value=mock_return) + with patch( + "litellm.proxy.management_helpers.user_invitation.create_invitation_for_user", + mock_create, + ): + response = client_with_auth.post(endpoint, json=payload) + else: + mock_prisma.db.litellm_invitationlink.find_unique = AsyncMock( + return_value={**mock_return, "created_by": "admin-user-id"} + ) + mock_prisma.db.litellm_invitationlink.delete = AsyncMock( + return_value=mock_return + ) + response = client_with_auth.post(endpoint, json=payload) + + assert response.status_code == 200 + data = response.json() + assert data["id"] == mock_return["id"] + assert data["user_id"] == mock_return["user_id"] + + @pytest.mark.parametrize( + "endpoint,payload", + [ + ("/invitation/new", {"user_id": "target-user-123"}), + ("/invitation/delete", {"invitation_id": "inv-456"}), + ], + ) + def test_invitation_endpoints_non_admin_denied( + self, client_with_auth, endpoint, payload + ): + """Non-admin users cannot access invitation endpoints.""" + from litellm.proxy._types import LitellmUserRoles + + mock_auth = MagicMock() + mock_auth.user_id = "regular-user" + mock_auth.user_role = LitellmUserRoles.INTERNAL_USER + mock_auth.api_key = "sk-regular" + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_invitationlink = MagicMock() + # Avoid triggering async DB calls in _user_has_admin_privileges + with patch( + "litellm.proxy.proxy_server._user_has_admin_privileges", + new_callable=AsyncMock, + return_value=False, + ): + response = client_with_auth.post(endpoint, json=payload) + + assert response.status_code == 400 + body = response.json() + # ProxyException handler returns {"error": {...}}, HTTPException returns {"detail": {...}} + error_content = body.get("error", body.get("detail", body)) + assert "not allowed" in str(error_content).lower() diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 5074bbf4397..6d6162437c4 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -7,6 +7,7 @@ sys.path.insert( from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, + TOOL_CALLS_CACHE, ) from litellm.types.llms.openai import ( ChatCompletionResponseMessage, @@ -17,6 +18,8 @@ from litellm.types.utils import ( CompletionTokensDetailsWrapper, Message, ModelResponse, + Function, + ChatCompletionMessageToolCall, PromptTokensDetailsWrapper, Usage, ) @@ -755,6 +758,98 @@ class TestFunctionCallTransformation: tool_call = tool_calls[0] assert tool_call.get("id") == "fallback_id" + def test_ensure_tool_results_preserves_cached_openai_object_tool_call(self): + """ + Test cached ChatCompletionMessageToolCall objects are normalized correctly. + """ + tool_call_id = "call_cached_openai_object" + TOOL_CALLS_CACHE.set_cache( + key=tool_call_id, + value=ChatCompletionMessageToolCall( + id=tool_call_id, + type="function", + function=Function( + name="search_web", + arguments='{"query": "python bugs"}', + ), + ), + ) + + messages_missing_tool_calls = [ + {"role": "user", "content": "Search for python bugs"}, + {"role": "assistant", "content": None, "tool_calls": []}, + {"role": "tool", "content": "Found 5 results", "tool_call_id": tool_call_id}, + ] + + try: + fixed_messages = LiteLLMCompletionResponsesConfig._ensure_tool_results_have_corresponding_tool_calls( + messages=messages_missing_tool_calls, + tools=None, + ) + finally: + TOOL_CALLS_CACHE.delete_cache(key=tool_call_id) + + assistant_msg = fixed_messages[1] + tool_calls = assistant_msg.get("tool_calls", []) + assert len(tool_calls) == 1 + + tool_call = tool_calls[0] + function = tool_call.get("function", {}) + assert function.get("name") == "search_web" + assert function.get("arguments") == '{"query": "python bugs"}' + + def test_ensure_tool_results_preserves_cached_attr_object_tool_call(self): + """ + Test cached attribute-only tool call objects are normalized correctly. + """ + + class AttrOnlyFunction: + def __init__(self, name: str, arguments: str): + self.name = name + self.arguments = arguments + + class AttrOnlyToolCall: + def __init__(self, id: str, type: str, function: AttrOnlyFunction): + self.id = id + self.type = type + self.function = function + + tool_call_id = "call_cached_attr_object" + TOOL_CALLS_CACHE.set_cache( + key=tool_call_id, + value=AttrOnlyToolCall( + id=tool_call_id, + type="function", + function=AttrOnlyFunction( + name="search_web", + arguments='{"query": "attribute objects"}', + ), + ), + ) + + messages_missing_tool_calls = [ + {"role": "user", "content": "Search using attr object"}, + {"role": "assistant", "content": None, "tool_calls": []}, + {"role": "tool", "content": "Found 3 results", "tool_call_id": tool_call_id}, + ] + + try: + fixed_messages = LiteLLMCompletionResponsesConfig._ensure_tool_results_have_corresponding_tool_calls( + messages=messages_missing_tool_calls, + tools=None, + ) + finally: + TOOL_CALLS_CACHE.delete_cache(key=tool_call_id) + + assistant_msg = fixed_messages[1] + tool_calls = assistant_msg.get("tool_calls", []) + assert len(tool_calls) == 1 + + tool_call = tool_calls[0] + function = tool_call.get("function", {}) + assert function.get("name") == "search_web" + assert function.get("arguments") == '{"query": "attribute objects"}' + class TestToolChoiceTransformation: """Test the tool_choice transformation fix for Cursor IDE bug""" @@ -1424,6 +1519,47 @@ class TestUsageTransformation: assert response_usage.input_tokens_details is None assert response_usage.output_tokens_details is None + def test_transform_usage_with_image_tokens(self): + """Test that image_tokens from Vertex AI/Gemini are properly transformed to output_tokens_details""" + # Setup: Simulate Vertex AI/Gemini usage with image_tokens in completion_tokens_details + usage = Usage( + prompt_tokens=10, + completion_tokens=150, + total_tokens=160, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=0, + text_tokens=50, + image_tokens=100, # From Vertex AI candidatesTokensDetails with modality="IMAGE" + ), + ) + + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="gemini-2.0-flash", + object="chat.completion", + usage=usage, + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="Here is the generated image.", role="assistant"), + ) + ], + ) + + # Execute + response_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + chat_completion_response=chat_completion_response + ) + + # Assert + assert response_usage.output_tokens == 150 + assert response_usage.output_tokens_details is not None + assert response_usage.output_tokens_details.reasoning_tokens == 0 + assert response_usage.output_tokens_details.text_tokens == 50 + assert response_usage.output_tokens_details.image_tokens == 100 + class TestStreamingIDConsistency: """Test cases for consistent IDs across streaming events (issue #14962)""" @@ -1637,4 +1773,4 @@ class TestStreamingIDConsistency: # Verify it matches the cached ID assert iterator._cached_item_id is not None - assert iterator._cached_item_id == text_done_id \ No newline at end of file + assert iterator._cached_item_id == text_done_id diff --git a/tests/test_litellm/secret_managers/test_secret_managers_main.py b/tests/test_litellm/secret_managers/test_secret_managers_main.py index eaef6956cd5..4a6e303586a 100644 --- a/tests/test_litellm/secret_managers/test_secret_managers_main.py +++ b/tests/test_litellm/secret_managers/test_secret_managers_main.py @@ -46,15 +46,24 @@ def mock_env(): yield os.environ -@patch("litellm.secret_managers.main.oidc_cache") -@patch("litellm.secret_managers.main._get_oidc_http_handler") -@patch("httpx.Client") # Prevent any real HTTP connections -def test_oidc_google_success(mock_httpx_client, mock_get_http_handler, mock_oidc_cache): - mock_oidc_cache.get_cache.return_value = None - mock_handler = MockHTTPHandler(timeout=600.0) - mock_get_http_handler.return_value = mock_handler +def test_oidc_google_success(): + """Test Google OIDC token fetch with mocked handler (no real network calls).""" secret_name = "oidc/google/[invalid url, do not cite]" - result = get_secret(secret_name) + mock_handler = MockHTTPHandler(timeout=600.0) + mock_get_http_handler = Mock(return_value=mock_handler) + mock_oidc_cache = Mock() + mock_oidc_cache.get_cache.return_value = None + + with patch("litellm.secret_managers.main.oidc_cache", mock_oidc_cache): + with patch( + "litellm.secret_managers.main._get_oidc_http_handler", + mock_get_http_handler, + ): + with patch( + "litellm.secret_managers.main.HTTPHandler", + side_effect=lambda timeout=None: mock_handler, + ): + result = get_secret(secret_name) assert result == "mocked_token" assert mock_handler.last_params == {"audience": "[invalid url, do not cite]"} @@ -63,32 +72,49 @@ def test_oidc_google_success(mock_httpx_client, mock_get_http_handler, mock_oidc ) -@patch("litellm.secret_managers.main.oidc_cache") -@patch("litellm.secret_managers.main._get_oidc_http_handler") -def test_oidc_google_cached(mock_get_http_handler, mock_oidc_cache): +def test_oidc_google_cached(): + """Test Google OIDC uses cache and does not call HTTP (no real network calls).""" + secret_name = "oidc/google/[invalid url, do not cite]" + mock_get_http_handler = Mock() + mock_oidc_cache = Mock() mock_oidc_cache.get_cache.return_value = "cached_token" - secret_name = "oidc/google/[invalid url, do not cite]" - result = get_secret(secret_name) + with patch("litellm.secret_managers.main.oidc_cache", mock_oidc_cache): + with patch( + "litellm.secret_managers.main._get_oidc_http_handler", + mock_get_http_handler, + ): + with patch( + "litellm.secret_managers.main.HTTPHandler", + Mock(side_effect=AssertionError("HTTPHandler should not be used")), + ): + result = get_secret(secret_name) assert result == "cached_token", f"Expected cached token, got {result}" mock_oidc_cache.get_cache.assert_called_with(key=secret_name) - # Verify HTTP handler was never called since we had a cached token mock_get_http_handler.assert_not_called() -@patch("litellm.secret_managers.main.oidc_cache") -@patch("litellm.secret_managers.main._get_oidc_http_handler") -def test_oidc_google_failure(mock_get_http_handler, mock_oidc_cache): +def test_oidc_google_failure(): + """Test Google OIDC raises when provider returns error (no real network calls).""" + secret_name = "oidc/google/https://example.com/api" mock_handler = MockHTTPHandler(timeout=600.0) mock_handler.status_code = 400 - mock_get_http_handler.return_value = mock_handler + mock_get_http_handler = Mock(return_value=mock_handler) + mock_oidc_cache = Mock() mock_oidc_cache.get_cache.return_value = None - - secret_name = "oidc/google/https://example.com/api" - with pytest.raises(ValueError, match="Google OIDC provider failed"): - get_secret(secret_name) + with patch("litellm.secret_managers.main.oidc_cache", mock_oidc_cache): + with patch( + "litellm.secret_managers.main._get_oidc_http_handler", + mock_get_http_handler, + ): + with patch( + "litellm.secret_managers.main.HTTPHandler", + side_effect=lambda timeout=None: mock_handler, + ): + with pytest.raises(ValueError, match="Google OIDC provider failed"): + get_secret(secret_name) def test_oidc_circleci_success(monkeypatch): @@ -151,20 +177,18 @@ def test_oidc_azure_file_success(mock_env, tmp_path): @patch("litellm.secret_managers.main.get_azure_ad_token_provider") -@patch.dict(os.environ, {}, clear=False) # Ensure AZURE_FEDERATED_TOKEN_FILE is not set -def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider): - # Ensure the env var is not set so it falls through to Azure AD token provider - if "AZURE_FEDERATED_TOKEN_FILE" in os.environ: - del os.environ["AZURE_FEDERATED_TOKEN_FILE"] - +def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider, monkeypatch): + # Force-unset so we always hit the Azure AD token provider path (CI may set AZURE_FEDERATED_TOKEN_FILE) + monkeypatch.delenv("AZURE_FEDERATED_TOKEN_FILE", raising=False) + # Mock the token provider function that gets returned and called mock_token_provider = Mock(return_value="azure_ad_token") mock_get_azure_ad_token_provider.return_value = mock_token_provider - + # Also mock the Azure Identity SDK to prevent any real Azure calls with patch("azure.identity.get_bearer_token_provider") as mock_bearer: mock_bearer.return_value = mock_token_provider - + secret_name = "oidc/azure/api://azure-audience" result = get_secret(secret_name) diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py new file mode 100644 index 00000000000..880e96f40a9 --- /dev/null +++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py @@ -0,0 +1,423 @@ +""" +Test suite for Anthropic beta headers filtering and mapping across all providers. + +This test validates: +1. Headers with null values in the config are filtered out +2. Headers with non-null values are correctly mapped to provider-specific names +3. Unknown headers (not in config) are filtered out +4. For Bedrock providers, beta headers appear in the request body (not just HTTP headers) +""" +import json +import os +from typing import Dict, List +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.anthropic_beta_headers_manager import ( + filter_and_transform_beta_headers, +) + + +class TestAnthropicBetaHeadersFiltering: + """Test beta header filtering and mapping for all providers.""" + + @pytest.fixture(autouse=True) + def setup(self): + """Load the beta headers config for testing.""" + config_path = os.path.join( + os.path.dirname(litellm.__file__), + "anthropic_beta_headers_config.json", + ) + with open(config_path, "r") as f: + self.config = json.load(f) + + def get_all_beta_headers(self) -> List[str]: + """Get all beta headers from the anthropic provider config.""" + return list(self.config.get("anthropic", {}).keys()) + + def get_supported_headers(self, provider: str) -> List[str]: + """Get headers with non-null values for a provider.""" + provider_config = self.config.get(provider, {}) + return [ + header for header, value in provider_config.items() if value is not None + ] + + def get_unsupported_headers(self, provider: str) -> List[str]: + """Get headers with null values for a provider.""" + provider_config = self.config.get(provider, {}) + return [header for header, value in provider_config.items() if value is None] + + def get_mapped_headers(self, provider: str) -> Dict[str, str]: + """Get mapping of input headers to provider-specific headers.""" + provider_config = self.config.get(provider, {}) + return { + header: value + for header, value in provider_config.items() + if value is not None + } + + @pytest.mark.parametrize( + "provider", + ["anthropic", "azure_ai", "bedrock_converse", "bedrock", "vertex_ai"], + ) + def test_filter_and_transform_beta_headers_all_headers(self, provider): + """Test filtering with all possible beta headers.""" + all_headers = self.get_all_beta_headers() + supported_headers = self.get_supported_headers(provider) + unsupported_headers = self.get_unsupported_headers(provider) + mapped_headers = self.get_mapped_headers(provider) + + filtered = filter_and_transform_beta_headers( + beta_headers=all_headers, provider=provider + ) + + for header in unsupported_headers: + assert ( + header not in filtered + ), f"Unsupported header '{header}' should be filtered out for {provider}" + assert ( + mapped_headers.get(header) not in filtered + ), f"Mapped value of unsupported header '{header}' should not appear for {provider}" + + for header in supported_headers: + expected_mapped = mapped_headers[header] + assert ( + expected_mapped in filtered + ), f"Supported header '{header}' should be mapped to '{expected_mapped}' for {provider}" + + @pytest.mark.parametrize( + "provider", + ["anthropic", "azure_ai", "bedrock_converse", "bedrock", "vertex_ai"], + ) + def test_unknown_headers_filtered_out(self, provider): + """Test that headers not in the config are filtered out.""" + unknown_headers = [ + "unknown-header-1", + "unknown-header-2", + "fake-beta-2025-01-01", + ] + all_headers = self.get_all_beta_headers() + unknown_headers + + filtered = filter_and_transform_beta_headers( + beta_headers=all_headers, provider=provider + ) + + for unknown in unknown_headers: + assert ( + unknown not in filtered + ), f"Unknown header '{unknown}' should be filtered out for {provider}" + + @pytest.mark.asyncio + async def test_anthropic_messages_http_headers_filtering(self): + """Test that Anthropic messages API filters HTTP headers correctly.""" + all_headers = self.get_all_beta_headers() + unsupported = self.get_unsupported_headers("anthropic") + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_client_factory: + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Hello"}], + "model": "claude-3-5-sonnet-20241022", + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + mock_response.headers = {} + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client_factory.return_value = mock_client + + try: + await litellm.acompletion( + model="anthropic/claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "Hi"}], + extra_headers={"anthropic-beta": ",".join(all_headers)}, + mock_response="Hello", + ) + except Exception: + pass + + if mock_client.post.called: + call_kwargs = mock_client.post.call_args.kwargs + headers = call_kwargs.get("headers", {}) + beta_header = headers.get("anthropic-beta", "") + + if beta_header: + beta_values = [b.strip() for b in beta_header.split(",")] + for unsupported_header in unsupported: + assert ( + unsupported_header not in beta_values + ), f"Unsupported header '{unsupported_header}' should not be in HTTP headers for Anthropic" + + @pytest.mark.asyncio + async def test_azure_ai_messages_http_headers_filtering(self): + """Test that Azure AI messages API filters HTTP headers correctly.""" + all_headers = self.get_all_beta_headers() + unsupported = self.get_unsupported_headers("azure_ai") + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_client_factory: + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Hello"}], + "model": "claude-3-5-sonnet-20241022", + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + mock_response.headers = {} + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client_factory.return_value = mock_client + + try: + await litellm.acompletion( + model="azure_ai/claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "Hi"}], + api_key="test-key", + api_base="https://test.azure.com", + extra_headers={"anthropic-beta": ",".join(all_headers)}, + mock_response="Hello", + ) + except Exception: + pass + + if mock_client.post.called: + call_kwargs = mock_client.post.call_args.kwargs + headers = call_kwargs.get("headers", {}) + beta_header = headers.get("anthropic-beta", "") + + if beta_header: + beta_values = [b.strip() for b in beta_header.split(",")] + for unsupported_header in unsupported: + assert ( + unsupported_header not in beta_values + ), f"Unsupported header '{unsupported_header}' should not be in HTTP headers for Azure AI" + + @pytest.mark.asyncio + async def test_bedrock_converse_headers_and_body_filtering(self): + """Test that Bedrock Converse filters both HTTP headers and request body correctly.""" + all_headers = self.get_all_beta_headers() + unsupported = self.get_unsupported_headers("bedrock_converse") + mapped_headers = self.get_mapped_headers("bedrock_converse") + + with patch("httpx.AsyncClient") as mock_client_class: + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "output": {"message": {"role": "assistant", "content": [{"text": "Hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 20}, + } + mock_response.headers = {} + mock_response.raise_for_status = MagicMock() + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client_class.return_value.__aenter__.return_value = mock_client + + try: + await litellm.acompletion( + model="bedrock/converse/us.anthropic.claude-3-5-sonnet-20241022-v2:0", + messages=[{"role": "user", "content": "Hi"}], + aws_access_key_id="test", + aws_secret_access_key="test", + aws_region_name="us-east-1", + extra_headers={"anthropic-beta": ",".join(all_headers)}, + mock_response="Hello", + ) + except Exception: + pass + + if mock_client.post.called: + call_kwargs = mock_client.post.call_args.kwargs + headers = call_kwargs.get("headers", {}) + beta_header = headers.get("anthropic-beta", "") + + if beta_header: + beta_values = [b.strip() for b in beta_header.split(",")] + for unsupported_header in unsupported: + assert ( + unsupported_header not in beta_values + ), f"Unsupported header '{unsupported_header}' should not be in HTTP headers for Bedrock Converse" + + data = call_kwargs.get("data") + if data: + body = json.loads(data) + body_beta = body.get("additionalModelRequestFields", {}).get( + "anthropic_beta", [] + ) + + for unsupported_header in unsupported: + assert ( + unsupported_header not in body_beta + ), f"Unsupported header '{unsupported_header}' should not be in request body for Bedrock Converse" + + for header, mapped_value in mapped_headers.items(): + if header in all_headers and mapped_value in body_beta: + assert ( + mapped_value in body_beta + ), f"Supported header '{header}' should be mapped to '{mapped_value}' in request body for Bedrock Converse" + + @pytest.mark.asyncio + async def test_vertex_ai_messages_http_headers_filtering(self): + """Test that Vertex AI messages API filters HTTP headers correctly.""" + all_headers = self.get_all_beta_headers() + unsupported = self.get_unsupported_headers("vertex_ai") + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_client_factory: + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Hello"}], + "model": "claude-3-5-sonnet-20241022", + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + mock_response.headers = {} + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client_factory.return_value = mock_client + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token" + ) as mock_token: + mock_token.return_value = ("test-token", "test-project") + + try: + await litellm.acompletion( + model="vertex_ai/claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "Hi"}], + vertex_project="test-project", + vertex_location="us-central1", + extra_headers={"anthropic-beta": ",".join(all_headers)}, + mock_response="Hello", + ) + except Exception: + pass + + if mock_client.post.called: + call_kwargs = mock_client.post.call_args.kwargs + headers = call_kwargs.get("headers", {}) + beta_header = headers.get("anthropic-beta", "") + + if beta_header: + beta_values = [b.strip() for b in beta_header.split(",")] + for unsupported_header in unsupported: + assert ( + unsupported_header not in beta_values + ), f"Unsupported header '{unsupported_header}' should not be in HTTP headers for Vertex AI" + + def test_header_mapping_correctness(self): + """Test that headers are mapped correctly for providers with transformations.""" + test_cases = [ + { + "provider": "bedrock", + "input": "advanced-tool-use-2025-11-20", + "expected": "tool-search-tool-2025-10-19", + }, + { + "provider": "vertex_ai", + "input": "advanced-tool-use-2025-11-20", + "expected": "tool-search-tool-2025-10-19", + }, + { + "provider": "anthropic", + "input": "advanced-tool-use-2025-11-20", + "expected": "advanced-tool-use-2025-11-20", + }, + { + "provider": "bedrock_converse", + "input": "computer-use-2025-01-24", + "expected": "computer-use-2025-01-24", + }, + { + "provider": "azure_ai", + "input": "advanced-tool-use-2025-11-20", + "expected": "advanced-tool-use-2025-11-20", + }, + ] + + for test_case in test_cases: + filtered = filter_and_transform_beta_headers( + beta_headers=[test_case["input"]], provider=test_case["provider"] + ) + + assert ( + test_case["expected"] in filtered + ), f"Header '{test_case['input']}' should be mapped to '{test_case['expected']}' for {test_case['provider']}, but got: {filtered}" + + def test_null_value_headers_filtered(self): + """Test that headers with null values are always filtered out.""" + for provider in ["anthropic", "azure_ai", "bedrock_converse", "bedrock", "vertex_ai"]: + unsupported = self.get_unsupported_headers(provider) + + if unsupported: + filtered = filter_and_transform_beta_headers( + beta_headers=unsupported, provider=provider + ) + + assert ( + len(filtered) == 0 + ), f"All null-value headers should be filtered out for {provider}, but got: {filtered}" + + def test_empty_headers_list(self): + """Test that empty headers list returns empty result.""" + for provider in ["anthropic", "azure_ai", "bedrock_converse", "bedrock", "vertex_ai"]: + filtered = filter_and_transform_beta_headers( + beta_headers=[], provider=provider + ) + + assert ( + len(filtered) == 0 + ), f"Empty headers list should return empty result for {provider}" + + def test_mixed_supported_and_unsupported_headers(self): + """Test filtering with a mix of supported, unsupported, and unknown headers.""" + for provider in ["anthropic", "azure_ai", "bedrock_converse", "bedrock", "vertex_ai"]: + supported = self.get_supported_headers(provider) + unsupported = self.get_unsupported_headers(provider) + mapped_headers = self.get_mapped_headers(provider) + + if not supported or not unsupported: + continue + + test_headers = ( + [supported[0]] + + [unsupported[0]] + + ["unknown-header-123"] + ) + + filtered = filter_and_transform_beta_headers( + beta_headers=test_headers, provider=provider + ) + + expected_mapped = mapped_headers[supported[0]] + assert ( + expected_mapped in filtered + ), f"Supported header should be in result for {provider}" + assert ( + unsupported[0] not in filtered + ), f"Unsupported header should not be in result for {provider}" + assert ( + "unknown-header-123" not in filtered + ), f"Unknown header should not be in result for {provider}" diff --git a/tests/test_litellm/test_anthropic_beta_headers_manager.py b/tests/test_litellm/test_anthropic_beta_headers_manager.py deleted file mode 100644 index d161426c22e..00000000000 --- a/tests/test_litellm/test_anthropic_beta_headers_manager.py +++ /dev/null @@ -1,306 +0,0 @@ -""" -Tests for the centralized Anthropic beta headers manager. - -Design: JSON config lists UNSUPPORTED headers for each provider. -Headers not in the unsupported list are passed through. -Header transformations (e.g., advanced-tool-use -> tool-search-tool) happen in code, not in JSON. -""" - -import pytest - -from litellm.anthropic_beta_headers_manager import ( - filter_and_transform_beta_headers, - get_provider_beta_header, - get_provider_name, - get_unsupported_headers, - is_beta_header_supported, - update_headers_with_filtered_beta, -) - - -class TestProviderNameResolution: - """Test provider name resolution and aliases.""" - - def test_get_provider_name_direct(self): - """Test direct provider names.""" - assert get_provider_name("anthropic") == "anthropic" - assert get_provider_name("bedrock") == "bedrock" - assert get_provider_name("vertex_ai") == "vertex_ai" - assert get_provider_name("azure_ai") == "azure_ai" - - def test_get_provider_name_alias(self): - """Test provider aliases.""" - # Note: Aliases are defined in the JSON config - # If no alias exists, the original name is returned - assert get_provider_name("azure") == "azure" # No alias defined - assert get_provider_name("vertex_ai_beta") == "vertex_ai_beta" # No alias defined - - -class TestBetaHeaderSupport: - """Test beta header support checks (unsupported list approach).""" - - def test_anthropic_supports_all_headers(self): - """Anthropic should support all beta headers (empty unsupported list).""" - headers = [ - "web-fetch-2025-09-10", - "web-search-2025-03-05", - "context-management-2025-06-27", - "compact-2026-01-12", - "structured-outputs-2025-11-13", - "advanced-tool-use-2025-11-20", - ] - for header in headers: - assert is_beta_header_supported(header, "anthropic") - - def test_bedrock_unsupported_headers(self): - """Bedrock should block specific headers.""" - # Not supported (in unsupported list) - assert not is_beta_header_supported("advanced-tool-use-2025-11-20", "bedrock") - assert not is_beta_header_supported( - "prompt-caching-scope-2026-01-05", "bedrock" - ) - assert not is_beta_header_supported("structured-outputs-2025-11-13", "bedrock") - - # Supported (not in unsupported list) - assert is_beta_header_supported("context-management-2025-06-27", "bedrock") - assert is_beta_header_supported("effort-2025-11-24", "bedrock") - assert is_beta_header_supported("tool-examples-2025-10-29", "bedrock") - - def test_vertex_ai_unsupported_headers(self): - """Vertex AI should block specific headers.""" - # Not supported (in unsupported list) - assert not is_beta_header_supported( - "prompt-caching-scope-2026-01-05", "vertex_ai" - ) - - # Supported (not in unsupported list) - assert is_beta_header_supported("web-search-2025-03-05", "vertex_ai") - assert is_beta_header_supported("context-management-2025-06-27", "vertex_ai") - assert is_beta_header_supported("effort-2025-11-24", "vertex_ai") - assert is_beta_header_supported("advanced-tool-use-2025-11-20", "vertex_ai") - - -class TestBetaHeaderTransformation: - """Test beta header support checking (transformations happen in code, not here).""" - - def test_anthropic_no_transformation(self): - """Anthropic headers should pass through (empty unsupported list).""" - header = "advanced-tool-use-2025-11-20" - assert get_provider_beta_header(header, "anthropic") == header - - def test_bedrock_unsupported_returns_none(self): - """Bedrock should return None for unsupported headers.""" - header = "advanced-tool-use-2025-11-20" - # This header is in bedrock's unsupported list - assert get_provider_beta_header(header, "bedrock") is None - - def test_vertex_ai_supported_returns_original(self): - """Vertex AI should return original for supported headers.""" - header = "advanced-tool-use-2025-11-20" - # This header is NOT in vertex_ai's unsupported list - assert get_provider_beta_header(header, "vertex_ai") == header - - def test_unsupported_header_returns_none(self): - """Unsupported headers (in unsupported list) should return None.""" - header = "prompt-caching-scope-2026-01-05" - assert get_provider_beta_header(header, "bedrock") is None - - def test_supported_header_returns_original(self): - """Supported headers (not in unsupported list) should return original.""" - header = "context-management-2025-06-27" - assert get_provider_beta_header(header, "bedrock") == header - - -class TestFilterAndTransformBetaHeaders: - """Test the main filtering and transformation function.""" - - def test_anthropic_keeps_all_headers(self): - """Anthropic should keep all headers (empty unsupported list).""" - headers = [ - "web-fetch-2025-09-10", - "context-management-2025-06-27", - "structured-outputs-2025-11-13", - "some-new-future-header-2026-01-01", # Even unknown headers pass through - ] - result = filter_and_transform_beta_headers(headers, "anthropic") - assert set(result) == set(headers) - - def test_bedrock_filters_unsupported(self): - """Bedrock should filter out headers in unsupported list.""" - headers = [ - "context-management-2025-06-27", # Not in unsupported list -> kept - "advanced-tool-use-2025-11-20", # In unsupported list -> dropped - "structured-outputs-2025-11-13", # In unsupported list -> dropped - "prompt-caching-scope-2026-01-05", # In unsupported list -> dropped - ] - result = filter_and_transform_beta_headers(headers, "bedrock") - assert "context-management-2025-06-27" in result - assert "advanced-tool-use-2025-11-20" not in result - assert "structured-outputs-2025-11-13" not in result - assert "prompt-caching-scope-2026-01-05" not in result - - def test_bedrock_no_transformations_in_filter(self): - """Bedrock filtering doesn't do transformations (those happen in code).""" - headers = ["advanced-tool-use-2025-11-20"] - result = filter_and_transform_beta_headers(headers, "bedrock") - # advanced-tool-use is in unsupported list, so it gets dropped - assert result == [] - - def test_vertex_ai_filters_unsupported(self): - """Vertex AI should filter unsupported headers.""" - headers = [ - "web-search-2025-03-05", # Not in unsupported list -> kept - "advanced-tool-use-2025-11-20", # Not in unsupported list -> kept - "prompt-caching-scope-2026-01-05", # In unsupported list -> dropped - ] - result = filter_and_transform_beta_headers(headers, "vertex_ai") - assert "web-search-2025-03-05" in result - assert "advanced-tool-use-2025-11-20" in result # Kept as-is, transformation happens in code - assert "prompt-caching-scope-2026-01-05" not in result - - def test_empty_list_returns_empty(self): - """Empty list should return empty list.""" - result = filter_and_transform_beta_headers([], "anthropic") - assert result == [] - - def test_bedrock_converse_more_restrictive(self): - """Bedrock Converse should be more restrictive than Bedrock.""" - headers = [ - "context-management-2025-06-27", - "advanced-tool-use-2025-11-20", - "tool-examples-2025-10-29", - ] - - bedrock_result = filter_and_transform_beta_headers(headers, "bedrock") - converse_result = filter_and_transform_beta_headers(headers, "bedrock_converse") - - # Bedrock Converse has more restrictions - # advanced-tool-use is in both unsupported lists - assert "advanced-tool-use-2025-11-20" not in bedrock_result - assert "advanced-tool-use-2025-11-20" not in converse_result - - # tool-examples is supported on bedrock but not converse - # Actually, looking at the JSON, tool-examples is NOT in bedrock unsupported list - # So it should be in bedrock_result - assert "tool-examples-2025-10-29" in bedrock_result - # But it's not explicitly in converse unsupported list either, so it passes through - # Let me check the actual behavior - assert "context-management-2025-06-27" in bedrock_result - assert "context-management-2025-06-27" in converse_result - - def test_unknown_future_headers_pass_through(self): - """Headers not in unsupported list should pass through (future-proof).""" - headers = ["some-new-beta-2026-05-01", "another-feature-2026-06-01"] - result = filter_and_transform_beta_headers(headers, "anthropic") - assert set(result) == set(headers) - - -class TestUpdateHeadersWithFilteredBeta: - """Test the headers update function.""" - - def test_update_headers_anthropic(self): - """Test updating headers for Anthropic.""" - headers = { - "anthropic-beta": "web-fetch-2025-09-10,context-management-2025-06-27" - } - result = update_headers_with_filtered_beta(headers, "anthropic") - assert "anthropic-beta" in result - beta_values = set(result["anthropic-beta"].split(",")) - assert "web-fetch-2025-09-10" in beta_values - assert "context-management-2025-06-27" in beta_values - - def test_update_headers_bedrock_filters(self): - """Test updating headers for Bedrock with filtering.""" - headers = { - "anthropic-beta": "context-management-2025-06-27,advanced-tool-use-2025-11-20" - } - result = update_headers_with_filtered_beta(headers, "bedrock") - assert "anthropic-beta" in result - assert "context-management-2025-06-27" in result["anthropic-beta"] - assert "advanced-tool-use-2025-11-20" not in result["anthropic-beta"] - - def test_update_headers_bedrock_no_transformations(self): - """Test that filtering doesn't do transformations (those happen in code).""" - headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} - result = update_headers_with_filtered_beta(headers, "bedrock") - # advanced-tool-use is in unsupported list, so it gets dropped - assert "anthropic-beta" not in result - - def test_update_headers_removes_if_all_filtered(self): - """Test that header is removed if all values are filtered.""" - headers = {"anthropic-beta": "advanced-tool-use-2025-11-20,prompt-caching-scope-2026-01-05"} - result = update_headers_with_filtered_beta(headers, "bedrock") - assert "anthropic-beta" not in result - - def test_update_headers_no_beta_header(self): - """Test updating headers when no beta header exists.""" - headers = {"content-type": "application/json"} - result = update_headers_with_filtered_beta(headers, "anthropic") - assert "anthropic-beta" not in result - assert headers == result - - -class TestGetUnsupportedHeaders: - """Test getting unsupported headers for a provider.""" - - def test_anthropic_has_no_unsupported(self): - """Anthropic should have no unsupported headers (empty list).""" - anthropic_unsupported = get_unsupported_headers("anthropic") - assert len(anthropic_unsupported) == 0 - - def test_bedrock_converse_most_restrictive(self): - """Bedrock Converse should have more unsupported headers than Bedrock.""" - bedrock_unsupported = get_unsupported_headers("bedrock") - converse_unsupported = get_unsupported_headers("bedrock_converse") - # Converse has more restrictions - assert len(converse_unsupported) >= len(bedrock_unsupported) - - def test_all_providers_have_config(self): - """All providers should have a configuration entry.""" - providers = ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai"] - for provider in providers: - unsupported = get_unsupported_headers(provider) - # Should return a list (even if empty) - assert isinstance(unsupported, list), f"Provider {provider} should return a list" - - -class TestEdgeCases: - """Test edge cases and error handling.""" - - def test_unknown_provider(self): - """Unknown provider with no config should pass through all headers.""" - result = filter_and_transform_beta_headers( - ["context-management-2025-06-27"], "unknown_provider" - ) - # Unknown providers have no unsupported list, so headers pass through - assert "context-management-2025-06-27" in result - - def test_whitespace_handling(self): - """Headers with whitespace should be handled correctly.""" - headers = [ - " context-management-2025-06-27 ", - " web-search-2025-03-05 ", - ] - result = filter_and_transform_beta_headers(headers, "anthropic") - assert len(result) == 2 - - def test_duplicate_headers(self): - """Duplicate headers should be deduplicated.""" - headers = [ - "context-management-2025-06-27", - "context-management-2025-06-27", - ] - result = filter_and_transform_beta_headers(headers, "anthropic") - assert len(result) == 1 - - def test_case_sensitivity(self): - """Headers should be case-sensitive.""" - # Correct case - should pass through for anthropic (no unsupported list) - headers = ["context-management-2025-06-27"] - result = filter_and_transform_beta_headers(headers, "anthropic") - assert len(result) == 1 - - # Wrong case - should still pass through (not in unsupported list) - headers = ["Context-Management-2025-06-27"] - result = filter_and_transform_beta_headers(headers, "anthropic") - assert len(result) == 1 # Passes through because anthropic has empty unsupported list diff --git a/tests/test_litellm/test_constants.py b/tests/test_litellm/test_constants.py index 77f2f308f88..23447a02e04 100644 --- a/tests/test_litellm/test_constants.py +++ b/tests/test_litellm/test_constants.py @@ -38,6 +38,11 @@ def test_all_numeric_constants_can_be_overridden(): print("all numeric constants", json.dumps(numeric_constants, indent=4)) + # Constants that use a different env var name than the constant name + constant_to_env_var = { + "MAX_CALLBACKS": "LITELLM_MAX_CALLBACKS", + } + # Verify all numeric constants have environment variable support for name, value in numeric_constants: # Skip constants that are not meant to be overridden (if any) @@ -47,8 +52,11 @@ def test_all_numeric_constants_can_be_overridden(): # Create a test value that's different from the default test_value = value + 1 if isinstance(value, int) else value + 0.1 + # Use the env var name that the constants module actually reads + env_var_name = constant_to_env_var.get(name, name) + # Set the environment variable - with mock.patch.dict(os.environ, {name: str(test_value)}): + with mock.patch.dict(os.environ, {env_var_name: str(test_value)}): print("overriding", name, "with", test_value) importlib.reload(constants) diff --git a/tests/test_litellm/test_deepseek_model_metadata.py b/tests/test_litellm/test_deepseek_model_metadata.py new file mode 100644 index 00000000000..4900af5d97d --- /dev/null +++ b/tests/test_litellm/test_deepseek_model_metadata.py @@ -0,0 +1,180 @@ +""" +Regression tests for #20885 – ``supports_response_schema`` (and related +capability flags) must be consistent between the bare model-name entry +(e.g. ``deepseek-chat``) and the provider-prefixed entry +(e.g. ``deepseek/deepseek-chat``) in the model-cost map. + +The bug caused ``supports_response_schema("deepseek/deepseek-chat")`` to +return ``False`` even though the canonical ``deepseek-chat`` entry has the +field set to ``True``. +""" + +import json +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.utils import ( + _supports_factory, + supports_response_schema, +) + + +# --------------------------------------------------------------------------- +# Data-level tests – verify the JSON files are in sync +# --------------------------------------------------------------------------- + + +def _load_backup_json() -> dict: + """Load the backup JSON directly from disk.""" + backup_path = os.path.join( + os.path.dirname(litellm.__file__), + "model_prices_and_context_window_backup.json", + ) + with open(backup_path, encoding="utf-8") as f: + return json.load(f) + + +class TestDeepSeekModelCostEntries: + """Verify that provider-prefixed DeepSeek entries contain the same + capability flags as their bare-name counterparts in the JSON files.""" + + def test_deepseek_chat_supports_response_schema_in_backup(self): + data = _load_backup_json() + entry = data.get("deepseek/deepseek-chat", {}) + assert entry.get("supports_response_schema") is True + + def test_deepseek_reasoner_supports_response_schema_in_backup(self): + data = _load_backup_json() + entry = data.get("deepseek/deepseek-reasoner", {}) + assert entry.get("supports_response_schema") is True + + def test_deepseek_chat_supports_system_messages_in_backup(self): + data = _load_backup_json() + entry = data.get("deepseek/deepseek-chat", {}) + assert entry.get("supports_system_messages") is True + + def test_deepseek_reasoner_supports_system_messages_in_backup(self): + data = _load_backup_json() + entry = data.get("deepseek/deepseek-reasoner", {}) + assert entry.get("supports_system_messages") is True + + def test_deepseek_chat_max_input_tokens_matches_bare_in_backup(self): + data = _load_backup_json() + bare = data.get("deepseek-chat", {}) + prefixed = data.get("deepseek/deepseek-chat", {}) + assert prefixed.get("max_input_tokens") == bare.get("max_input_tokens") + + def test_deepseek_reasoner_max_output_tokens_matches_bare_in_backup(self): + data = _load_backup_json() + bare = data.get("deepseek-reasoner", {}) + prefixed = data.get("deepseek/deepseek-reasoner", {}) + assert prefixed.get("max_output_tokens") == bare.get("max_output_tokens") + + def test_main_json_deepseek_chat_supports_response_schema(self): + main_path = os.path.join( + os.path.dirname(os.path.dirname(litellm.__file__)), + "model_prices_and_context_window.json", + ) + with open(main_path, encoding="utf-8") as f: + data = json.load(f) + entry = data.get("deepseek/deepseek-chat", {}) + assert entry.get("supports_response_schema") is True + + def test_main_json_deepseek_reasoner_supports_response_schema(self): + main_path = os.path.join( + os.path.dirname(os.path.dirname(litellm.__file__)), + "model_prices_and_context_window.json", + ) + with open(main_path, encoding="utf-8") as f: + data = json.load(f) + entry = data.get("deepseek/deepseek-reasoner", {}) + assert entry.get("supports_response_schema") is True + + +# --------------------------------------------------------------------------- +# API-level tests – verify supports_response_schema returns True +# --------------------------------------------------------------------------- + + +class TestSupportsResponseSchemaDeepSeek: + """All calling conventions for DeepSeek should return True for + ``supports_response_schema``.""" + + def test_provider_slash_model(self): + assert supports_response_schema(model="deepseek/deepseek-chat") is True + + def test_explicit_provider(self): + assert ( + supports_response_schema( + model="deepseek-chat", custom_llm_provider="deepseek" + ) + is True + ) + + def test_reasoner_provider_slash_model(self): + assert supports_response_schema(model="deepseek/deepseek-reasoner") is True + + def test_reasoner_explicit_provider(self): + assert ( + supports_response_schema( + model="deepseek-reasoner", custom_llm_provider="deepseek" + ) + is True + ) + + +# --------------------------------------------------------------------------- +# Fallback-logic test – bare model entry used when prefixed is incomplete +# --------------------------------------------------------------------------- + + +class TestBareModelFallback: + """When a provider-prefixed entry is missing a capability flag, the + ``_supports_factory`` fallback should consult the bare model-name + entry in ``litellm.model_cost``.""" + + def test_fallback_uses_bare_entry(self): + """Temporarily remove ``supports_response_schema`` from the prefixed + entry and verify the fallback still returns True.""" + key = "deepseek/deepseek-chat" + original = litellm.model_cost.get(key, {}).get("supports_response_schema") + try: + # Simulate the pre-fix state: field missing from prefixed entry + if key in litellm.model_cost: + litellm.model_cost[key].pop("supports_response_schema", None) + result = _supports_factory( + model="deepseek-chat", + custom_llm_provider="deepseek", + key="supports_response_schema", + ) + assert result is True + finally: + # Restore + if key in litellm.model_cost and original is not None: + litellm.model_cost[key]["supports_response_schema"] = original + + def test_no_fallback_when_explicitly_false(self): + """If the prefixed entry explicitly sets a capability to ``False``, + the fallback must NOT override it.""" + key = "deepseek/deepseek-reasoner" + # After the data fix, deepseek/deepseek-reasoner has + # supports_function_calling=false (matching the bare entry). + # Explicitly set it to False to test the guard. + original = litellm.model_cost.get(key, {}).get("supports_function_calling") + try: + if key in litellm.model_cost: + litellm.model_cost[key]["supports_function_calling"] = False + result = _supports_factory( + model="deepseek-reasoner", + custom_llm_provider="deepseek", + key="supports_function_calling", + ) + assert result is False + finally: + if key in litellm.model_cost and original is not None: + litellm.model_cost[key]["supports_function_calling"] = original diff --git a/tests/test_litellm/test_exception_exports.py b/tests/test_litellm/test_exception_exports.py new file mode 100644 index 00000000000..cde26295bad --- /dev/null +++ b/tests/test_litellm/test_exception_exports.py @@ -0,0 +1,31 @@ +""" +Test that all standard HTTP error exceptions are exported from litellm.__init__. +""" + +import litellm + + +def test_permission_denied_error_is_exported(): + """PermissionDeniedError (403) should be accessible as litellm.PermissionDeniedError.""" + assert hasattr(litellm, "PermissionDeniedError") + assert litellm.PermissionDeniedError is not None + + +def test_all_http_error_exceptions_exported(): + """All standard HTTP error exceptions should be accessible at module level.""" + expected_exceptions = [ + "BadRequestError", # 400 + "AuthenticationError", # 401 + "PermissionDeniedError", # 403 + "NotFoundError", # 404 + "Timeout", # 408 + "UnprocessableEntityError", # 422 + "RateLimitError", # 429 + "InternalServerError", # 500 + "BadGatewayError", # 502 + "ServiceUnavailableError", # 503 + ] + for exc_name in expected_exceptions: + assert hasattr(litellm, exc_name), ( + f"litellm.{exc_name} is not exported from litellm.__init__" + ) diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 7e5931d8c0f..6f65ada7459 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -1,27 +1,21 @@ import asyncio -import datetime import json import os import sys -import unittest -from typing import List, Optional, Tuple -from unittest.mock import ANY, MagicMock, Mock, patch +from typing import List -import httpx import pytest sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system-path -import io import logging import sys -import unittest -from contextlib import redirect_stdout import litellm from litellm._logging import ( ALL_LOGGERS, + JsonFormatter, _initialize_loggers_with_handler, _turn_on_json, verbose_logger, @@ -72,6 +66,117 @@ def test_json_mode_emits_one_record_per_logger(capfd): assert "timestamp" in obj, "`timestamp` key missing" +def test_json_formatter_parses_embedded_json_message(): + """ + Test that JsonFormatter parses embedded JSON in the message field and promotes + sub-fields to first-class JSON properties for downstream querying. + """ + formatter = JsonFormatter() + record = logging.LogRecord( + name="LiteLLM", + level=logging.DEBUG, + pathname="", + lineno=0, + msg='{"event": "giveup", "exception": "Connection failed", "model_name": "gpt-4"}', + args=(), + exc_info=None, + ) + output = formatter.format(record) + obj = json.loads(output) + # Standard fields preserved + assert "message" in obj + assert obj["level"] == "DEBUG" + assert "timestamp" in obj + # Embedded JSON fields promoted to top-level for querying + assert obj["event"] == "giveup" + assert obj["exception"] == "Connection failed" + assert obj["model_name"] == "gpt-4" + + +def test_json_formatter_includes_extra_attributes(): + """ + Test that JsonFormatter includes extra attributes from logger.debug("msg", extra={...}). + """ + formatter = JsonFormatter() + record = logging.LogRecord( + name="LiteLLM", + level=logging.DEBUG, + pathname="", + lineno=0, + msg="POST Request Sent from LiteLLM", + args=(), + exc_info=None, + ) + record.api_base = "https://api.openai.com" + record.authorization = "Bearer sk-***" + output = formatter.format(record) + obj = json.loads(output) + assert obj["message"] == "POST Request Sent from LiteLLM" + assert obj["api_base"] == "https://api.openai.com" + assert obj["authorization"] == "Bearer sk-***" + + +def test_json_formatter_plain_message_unchanged(): + """ + Test that non-JSON messages are passed through as-is in the message field. + """ + formatter = JsonFormatter() + record = logging.LogRecord( + name="LiteLLM", + level=logging.INFO, + pathname="", + lineno=0, + msg="Cache hit!", + args=(), + exc_info=None, + ) + output = formatter.format(record) + obj = json.loads(output) + assert obj["message"] == "Cache hit!" + assert "event" not in obj + assert "exception" not in obj + + +def test_json_formatter_parses_embedded_python_dict_repr(): + """ + Test that JsonFormatter parses Python dict repr (str/deployment) embedded in + plain text, e.g. from get_available_deployment logs. + Reproduces Roni's reported case. + """ + formatter = JsonFormatter() + msg = ( + "get_available_deployment for model: text-embedding-3-large, " + "Selected deployment: {'model_name': 'text-embedding-3-large', " + "'litellm_params': {'api_key': 'sk**********', 'tpm': 1000000, 'rpm': 2000, " + "'use_in_pass_through': False, 'use_litellm_proxy': False, " + "'merge_reasoning_content_in_choices': False, 'model': 'text-embedding-3-large'}, " + "'model_info': {'id': 'a624b057aec64ada48311', 'db_model': False}} " + "for model: text-embedding-3-large" + ) + record = logging.LogRecord( + name="LiteLLM Router", + level=logging.INFO, + pathname="", + lineno=0, + msg=msg, + args=(), + exc_info=None, + ) + output = formatter.format(record) + obj = json.loads(output) + assert "message" in obj + assert obj["level"] == "INFO" + # Python dict parsed and promoted to first-class properties + assert obj["model_name"] == "text-embedding-3-large" + assert "litellm_params" in obj + assert obj["litellm_params"]["api_key"] == "sk**********" + assert obj["litellm_params"]["tpm"] == 1000000 + assert obj["litellm_params"]["use_in_pass_through"] is False + assert "model_info" in obj + assert obj["model_info"]["id"] == "a624b057aec64ada48311" + assert obj["model_info"]["db_model"] is False + + def test_initialize_loggers_with_handler_sets_propagate_false(): """ Test that the initialize_loggers_with_handler function sets propagate to False for all loggers @@ -96,7 +201,7 @@ async def test_cache_hit_includes_custom_llm_provider(): test_custom_logger = CacheHitCustomLogger() original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] litellm.callbacks = [test_custom_logger] - + try: # First call - should be a cache miss response1 = await litellm.acompletion( @@ -105,10 +210,10 @@ async def test_cache_hit_includes_custom_llm_provider(): mock_response="test response", caching=True, ) - + # Wait for logging to complete await asyncio.sleep(0.5) - + # Second identical call - should be a cache hit response2 = await litellm.acompletion( model="gpt-3.5-turbo", @@ -116,38 +221,43 @@ async def test_cache_hit_includes_custom_llm_provider(): mock_response="test response", caching=True, ) - + # Wait for logging to complete await asyncio.sleep(0.5) - + # Verify we have logged events - assert len(test_custom_logger.logged_standard_logging_payloads) >= 2, \ - f"Expected at least 2 logged events, got {len(test_custom_logger.logged_standard_logging_payloads)}" - + assert ( + len(test_custom_logger.logged_standard_logging_payloads) >= 2 + ), f"Expected at least 2 logged events, got {len(test_custom_logger.logged_standard_logging_payloads)}" + # Find the cache hit event (should be the second call) cache_hit_payload = None for payload in test_custom_logger.logged_standard_logging_payloads: if payload.get("cache_hit") is True: cache_hit_payload = payload break - + # Verify cache hit event was found - assert cache_hit_payload is not None, "No cache hit event found in logged payloads" - + assert ( + cache_hit_payload is not None + ), "No cache hit event found in logged payloads" + # Verify custom_llm_provider is included in the cache hit payload - assert "custom_llm_provider" in cache_hit_payload, \ - "custom_llm_provider missing from cache hit standard logging payload" - + assert ( + "custom_llm_provider" in cache_hit_payload + ), "custom_llm_provider missing from cache hit standard logging payload" + # Verify custom_llm_provider has a valid value (should be "openai" for gpt-3.5-turbo) custom_llm_provider = cache_hit_payload["custom_llm_provider"] - assert custom_llm_provider is not None and custom_llm_provider != "", \ - f"custom_llm_provider should not be None or empty, got: {custom_llm_provider}" - + assert ( + custom_llm_provider is not None and custom_llm_provider != "" + ), f"custom_llm_provider should not be None or empty, got: {custom_llm_provider}" + print( f"Cache hit standard logging payload with custom_llm_provider: {custom_llm_provider}", json.dumps(cache_hit_payload, indent=2), ) - + finally: # Clean up litellm.callbacks = original_callbacks diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 75ec806ee17..9dcb16b545e 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1990,3 +1990,94 @@ async def test_anthropic_messages_call_type_is_cached(): # This assertion will FAIL if anthropic_messages is filtered out assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" + + +def test_update_kwargs_with_deployment_propagates_model_tags(): + """ + Test that deployment-level tags from litellm_params are merged into + kwargs metadata when _update_kwargs_with_deployment is called. + + This ensures model-level tags defined in config.yaml appear in SpendLogs. + See: https://github.com/BerriAI/litellm/issues/XXXX + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key", + "tags": ["openai-account", "production"], + }, + }, + ], + ) + + kwargs: dict = {"metadata": {}} + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + + # Deployment tags should be propagated to kwargs metadata + assert "tags" in kwargs["metadata"] + assert "openai-account" in kwargs["metadata"]["tags"] + assert "production" in kwargs["metadata"]["tags"] + + +def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): + """ + Test that when both request-level and deployment-level tags exist, + they are merged without duplicates. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key", + "tags": ["openai-account", "shared-tag"], + }, + }, + ], + ) + + # Simulate request that already has tags (from request body or key/team level) + kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + + # Both sources should be merged, no duplicates + assert "user-tag" in kwargs["metadata"]["tags"] + assert "openai-account" in kwargs["metadata"]["tags"] + assert "shared-tag" in kwargs["metadata"]["tags"] + assert kwargs["metadata"]["tags"].count("shared-tag") == 1 + + +def test_update_kwargs_with_deployment_no_tags(): + """ + Test that when deployment has no tags, kwargs metadata is not affected. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key", + }, + }, + ], + ) + + kwargs: dict = {"metadata": {}} + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + + # No tags key should be added if deployment has no tags + assert "tags" not in kwargs["metadata"] diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py new file mode 100644 index 00000000000..2112295e040 --- /dev/null +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -0,0 +1,264 @@ +""" +Test that per-deployment custom pricing does not pollute the shared backend +model key in litellm.model_cost. + +When two deployments share the same backend model (e.g. vertex_ai/gemini-2.5-flash) +and one has explicit zero-cost pricing in model_info, the other deployment +should still use the built-in pricing. +""" + +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm import Router + + +def test_should_not_pollute_shared_key_with_zero_cost_pricing(): + """ + When deployment A has input_cost_per_token=0 and deployment B has no + custom pricing, deployment B should still report the built-in pricing + (not zero). + """ + backend_model = "vertex_ai/gemini-2.5-flash" + + # Grab built-in pricing before creating any router + builtin_info = litellm.get_model_info(model=backend_model) + builtin_input_cost = builtin_info["input_cost_per_token"] + builtin_output_cost = builtin_info["output_cost_per_token"] + + # Sanity: built-in pricing should be non-zero for this model + assert builtin_input_cost > 0, "Test requires a model with non-zero built-in pricing" + assert builtin_output_cost > 0, "Test requires a model with non-zero built-in pricing" + + router = Router( + model_list=[ + # Deployment A: explicit zero-cost pricing + { + "model_name": "custom-zero-cost-model", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-1", + }, + "model_info": { + "id": "deployment-a-zero-cost", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + # Deployment B: no custom pricing, relies on built-in + { + "model_name": "standard-cost-model", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-2", + }, + "model_info": { + "id": "deployment-b-builtin-cost", + }, + }, + ], + ) + + # Deployment A: should report zero pricing via its unique model_id + info_a = router.get_deployment_model_info( + model_id="deployment-a-zero-cost", + model_name=backend_model, + ) + assert info_a is not None + assert info_a["input_cost_per_token"] == 0.0 + assert info_a["output_cost_per_token"] == 0.0 + + # Deployment B: should report built-in pricing, NOT zero + info_b = router.get_deployment_model_info( + model_id="deployment-b-builtin-cost", + model_name=backend_model, + ) + assert info_b is not None + assert info_b["input_cost_per_token"] == builtin_input_cost, ( + f"Deployment B should use built-in input cost {builtin_input_cost}, " + f"got {info_b['input_cost_per_token']}" + ) + assert info_b["output_cost_per_token"] == builtin_output_cost, ( + f"Deployment B should use built-in output cost {builtin_output_cost}, " + f"got {info_b['output_cost_per_token']}" + ) + + +def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): + """ + A deployment with custom (non-zero) pricing should not overwrite + the shared backend key's built-in pricing. + """ + backend_model = "vertex_ai/gemini-2.5-flash" + + builtin_info = litellm.get_model_info(model=backend_model) + builtin_input_cost = builtin_info["input_cost_per_token"] + + router = Router( + model_list=[ + # Deployment with custom high pricing + { + "model_name": "expensive-model", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-3", + }, + "model_info": { + "id": "deployment-expensive", + "input_cost_per_token": 0.99, + "output_cost_per_token": 0.99, + }, + }, + # Deployment relying on built-in pricing + { + "model_name": "standard-model", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-4", + }, + "model_info": { + "id": "deployment-standard", + }, + }, + ], + ) + + # Custom pricing deployment should see its custom values + info_expensive = router.get_deployment_model_info( + model_id="deployment-expensive", + model_name=backend_model, + ) + assert info_expensive is not None + assert info_expensive["input_cost_per_token"] == 0.99 + assert info_expensive["output_cost_per_token"] == 0.99 + + # Standard deployment should still see built-in pricing + info_standard = router.get_deployment_model_info( + model_id="deployment-standard", + model_name=backend_model, + ) + assert info_standard is not None + assert info_standard["input_cost_per_token"] == builtin_input_cost, ( + f"Standard deployment should use built-in pricing {builtin_input_cost}, " + f"got {info_standard['input_cost_per_token']}" + ) + + +def test_should_store_full_pricing_under_deployment_model_id(): + """ + Per-deployment pricing (including zero) should be stored and + retrievable via the unique model_id key in litellm.model_cost. + """ + backend_model = "vertex_ai/gemini-2.5-flash" + + router = Router( + model_list=[ + { + "model_name": "zero-cost-model", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-5", + }, + "model_info": { + "id": "deployment-zero-check", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + ], + ) + + # The model_id entry should exist and have the zero pricing + entry = litellm.model_cost.get("deployment-zero-check") + assert entry is not None, "Deployment should be registered by model_id" + assert entry["input_cost_per_token"] == 0.0 + assert entry["output_cost_per_token"] == 0.0 + + +def test_should_preserve_builtin_pricing_regardless_of_deployment_order(): + """ + The built-in pricing should be preserved no matter which deployment + is processed first (zero-cost first, or standard first). + """ + backend_model = "vertex_ai/gemini-2.5-flash" + + builtin_info = litellm.get_model_info(model=backend_model) + builtin_input_cost = builtin_info["input_cost_per_token"] + builtin_output_cost = builtin_info["output_cost_per_token"] + + # Order 1: standard first, then zero-cost + router1 = Router( + model_list=[ + { + "model_name": "standard-first", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-6", + }, + "model_info": {"id": "order1-standard"}, + }, + { + "model_name": "zero-cost-second", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-7", + }, + "model_info": { + "id": "order1-zero", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + ], + ) + + info_std_1 = router1.get_deployment_model_info( + model_id="order1-standard", model_name=backend_model + ) + assert info_std_1["input_cost_per_token"] == builtin_input_cost + assert info_std_1["output_cost_per_token"] == builtin_output_cost + + # Order 2: zero-cost first, then standard + router2 = Router( + model_list=[ + { + "model_name": "zero-cost-first", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-8", + }, + "model_info": { + "id": "order2-zero", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + { + "model_name": "standard-second", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-9", + }, + "model_info": {"id": "order2-standard"}, + }, + ], + ) + + info_std_2 = router2.get_deployment_model_info( + model_id="order2-standard", model_name=backend_model + ) + assert info_std_2["input_cost_per_token"] == builtin_input_cost, ( + f"Order should not matter. Expected {builtin_input_cost}, " + f"got {info_std_2['input_cost_per_token']}" + ) + assert info_std_2["output_cost_per_token"] == builtin_output_cost, ( + f"Order should not matter. Expected {builtin_output_cost}, " + f"got {info_std_2['output_cost_per_token']}" + ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 352125d16c4..794b3b87187 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -661,6 +661,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_url_context": {"type": "boolean"}, "supports_reasoning": {"type": "boolean"}, "supports_service_tier": {"type": "boolean"}, + "supports_preset": {"type": "boolean"}, "tool_use_system_prompt_tokens": {"type": "number"}, "tpm": {"type": "number"}, "supported_endpoints": { diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 5446a0a7b3f..c8cc292519b 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -798,7 +798,7 @@ def test_openai_transform_video_content_request_empty_params(): def test_video_content_handler_uses_get_for_openai(): """HTTP handler must use GET (not POST) for OpenAI content download.""" from litellm.types.router import GenericLiteLLMParams - + handler = BaseLLMHTTPHandler() config = OpenAIVideoConfig() @@ -807,7 +807,12 @@ def test_video_content_handler_uses_get_for_openai(): mock_response.content = b"mp4-bytes" mock_client.get.return_value = mock_response + # Patch both where _get_httpx_client is used and where it is defined so the mock + # is used regardless of import order / CI environment with patch( + "litellm.llms.custom_httpx.http_handler._get_httpx_client", + return_value=mock_client, + ), patch( "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", return_value=mock_client, ): diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/ui/litellm-dashboard/e2e_tests/constants.ts index b07bd68fcf1..58b56af0a2b 100644 --- a/ui/litellm-dashboard/e2e_tests/constants.ts +++ b/ui/litellm-dashboard/e2e_tests/constants.ts @@ -1 +1,6 @@ export const ADMIN_STORAGE_PATH = "admin.storageState.json"; + +export const E2E_UPDATE_LIMITS_KEY_ID_PREFIX = "102c"; +export const E2E_DELETE_KEY_ID_PREFIX = "94a5"; +export const E2E_DELETE_KEY_NAME = "e2eDeleteKey"; +export const E2E_REGENERATE_KEY_ID_PREFIX = "593a"; diff --git a/ui/litellm-dashboard/e2e_tests/tests/keys/deleteKey.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/keys/deleteKey.spec.ts new file mode 100644 index 00000000000..a5841316251 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/keys/deleteKey.spec.ts @@ -0,0 +1,25 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_DELETE_KEY_ID_PREFIX, E2E_DELETE_KEY_NAME } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; + +test.describe("Delete Key", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Able to delete a key", async ({ page }) => { + await navigateToPage(page, Page.ApiKeys); + await expect(page.getByRole("button", { name: "Next" })).toBeVisible(); + await page + .locator("button", { + hasText: E2E_DELETE_KEY_ID_PREFIX, + }) + .click(); + await page.getByRole("button", { name: "Delete Key" }).click(); + await page.getByRole("textbox", { name: E2E_DELETE_KEY_NAME }).click(); + await page.getByRole("textbox", { name: E2E_DELETE_KEY_NAME }).fill(E2E_DELETE_KEY_NAME); + const deleteButton = page.getByRole("button", { name: "Delete", exact: true }); + await expect(deleteButton).toBeEnabled(); + await deleteButton.click(); + await expect(page.getByText("Key deleted successfully")).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/keys/regenerateKey.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/keys/regenerateKey.spec.ts new file mode 100644 index 00000000000..0188a4f81ce --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/keys/regenerateKey.spec.ts @@ -0,0 +1,21 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_REGENERATE_KEY_ID_PREFIX } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; + +test.describe("Regenerate Key", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Able to regenerate a key", async ({ page }) => { + await navigateToPage(page, Page.ApiKeys); + await expect(page.getByRole("button", { name: "Next" })).toBeVisible(); + await page + .locator("button", { + hasText: E2E_REGENERATE_KEY_ID_PREFIX, + }) + .click(); + await page.getByRole("button", { name: "Regenerate Key" }).click(); + await page.getByRole("button", { name: "Regenerate", exact: true }).click(); + await expect(page.getByText("Virtual Key regenerated")).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/keys/updateKeyLimits.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/keys/updateKeyLimits.spec.ts new file mode 100644 index 00000000000..6cae36272ab --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/keys/updateKeyLimits.spec.ts @@ -0,0 +1,27 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_UPDATE_LIMITS_KEY_ID_PREFIX } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; + +test.describe("Update Key TPM and RPM Limits", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Able to update a key's TPM and RPM limits", async ({ page }) => { + await navigateToPage(page, Page.ApiKeys); + await expect(page.getByRole("button", { name: "Next" })).toBeVisible(); + await page + .locator("button", { + hasText: E2E_UPDATE_LIMITS_KEY_ID_PREFIX, + }) + .click(); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await page.getByRole("spinbutton", { name: "TPM Limit" }).click(); + await page.getByRole("spinbutton", { name: "TPM Limit" }).fill("123"); + await page.getByRole("spinbutton", { name: "RPM Limit" }).click(); + await page.getByRole("spinbutton", { name: "RPM Limit" }).fill("456"); + await page.getByRole("button", { name: "Save Changes" }).click(); + await expect(page.getByRole("paragraph").filter({ hasText: "TPM: 123" })).toBeVisible(); + await expect(page.getByRole("paragraph").filter({ hasText: "RPM: 456" })).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 4205657ca8c..3a21813fbf4 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -13159,6 +13159,21 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } + }, + "node_modules/@next/swc-win32-ia32-msvc": { + "version": "14.2.33", + "resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz", + "integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==", + "cpu": [ + "ia32" + ], + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + } } } } diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 76ac97f008c..164368eb6ba 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -3,7 +3,7 @@ "version": "0.1.0", "private": true, "scripts": { - "dev": "next dev --webpack", + "dev": "next dev", "build": "next build", "start": "next start", "lint": "next lint", @@ -79,6 +79,7 @@ "vitest": "^3.2.4" }, "overrides": { + "diff": ">=8.0.3", "prismjs": ">=1.30.0", "webpack-dev-server": ">=5.2.1", "mermaid": ">=11.10.0", diff --git a/ui/litellm-dashboard/scripts/e2e_tests/neonHelperScripts.ts b/ui/litellm-dashboard/scripts/e2e_tests/neonHelperScripts.ts index 3078a0d90d2..089ad4e7926 100644 --- a/ui/litellm-dashboard/scripts/e2e_tests/neonHelperScripts.ts +++ b/ui/litellm-dashboard/scripts/e2e_tests/neonHelperScripts.ts @@ -1,4 +1,4 @@ -import { createApiClient } from "@neondatabase/api-client"; +import { createApiClient, EndpointType } from "@neondatabase/api-client"; import { config } from "dotenv"; import { resolve } from "path"; @@ -27,6 +27,13 @@ export async function createNeonE2ETestingBranch(projectId: string, parentBranch parent_id: parentBranchId, expires_at: expireAt ?? new Date(Date.now() + 1000 * 60 * 30).toISOString(), }, + endpoints: [ + { + type: EndpointType.ReadWrite, + autoscaling_limit_min_cu: 0.25, + autoscaling_limit_max_cu: 1, + }, + ], }); return response; } catch (error) { @@ -35,13 +42,15 @@ export async function createNeonE2ETestingBranch(projectId: string, parentBranch } export async function getNeonE2ETestingBranchConnectionString() { - await createNeonE2ETestingBranch(PROJECT_ID, PARENT_BRANCH); - + const createBranchResponse = await createNeonE2ETestingBranch(PROJECT_ID, PARENT_BRANCH); + const projectId = createBranchResponse.data.branch.project_id; const response = await apiClient.getConnectionUri({ database_name: NEON_E2E_UI_TEST_DB_NAME, role_name: "neondb_owner", - projectId: PROJECT_ID, + projectId: projectId, }); console.log("connection string:", response.data.uri); return response.data.uri; } + +getNeonE2ETestingBranchConnectionString(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx index 405f8329b67..a74d3c108d6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx @@ -31,7 +31,7 @@ import { import * as React from "react"; import { useRouter, usePathname } from "next/navigation"; import { all_admin_roles, internalUserRoles, isAdminRole, rolesWithWriteAccess } from "@/utils/roles"; -import UsageIndicator from "@/components/usage_indicator"; +import UsageIndicator from "@/components/UsageIndicator"; import { serverRootPath } from "@/components/networking"; const { Sider } = Layout; @@ -64,7 +64,7 @@ const getBasePath = () => { const raw = process.env.NEXT_PUBLIC_BASE_URL ?? ""; const trimmed = raw.replace(/^\/+|\/+$/g, ""); // strip leading/trailing slashes const uiPath = trimmed ? `/${trimmed}/` : "/"; - + // If serverRootPath is set and not "/", prepend it to the UI path if (serverRootPath && serverRootPath !== "/") { // Remove trailing slash from serverRootPath and ensure uiPath has no leading slash for proper joining @@ -72,7 +72,7 @@ const getBasePath = () => { const cleanUiPath = uiPath.replace(/^\/+/, ""); return `${cleanServerRoot}/${cleanUiPath}`; } - + return uiPath; }; @@ -153,170 +153,170 @@ const toHref = (slugOrPath: string) => { // ----- Menu config (unchanged labels/icons; same appearance) ----- const menuItems: MenuItemCfg[] = [ - { key: "1", page: "api-keys", label: "Virtual Keys", icon: }, - { - key: "3", - page: "llm-playground", - label: "Test Key", - icon: , - roles: rolesWithWriteAccess, - }, - { - key: "2", - page: "models", - label: "Models + Endpoints", - icon: , - roles: rolesWithWriteAccess, - }, - { - key: "12", - page: "new_usage", - label: "Usage", - icon: , - roles: [...all_admin_roles, ...internalUserRoles], - }, - { key: "6", page: "teams", label: "Teams", icon: }, - { - key: "17", - page: "organizations", - label: "Organizations", - icon: , - roles: all_admin_roles, - }, - { - key: "5", - page: "users", - label: "Internal Users", - icon: , - roles: all_admin_roles, - }, - { key: "14", page: "api_ref", label: "API Reference", icon: }, - { - key: "16", - page: "model-hub-table", - label: "Model Hub", - icon: , - }, - { key: "15", page: "logs", label: "Logs", icon: }, - { - key: "11", - page: "guardrails", - label: "Guardrails", - icon: , - roles: all_admin_roles, - }, - { - key: "28", - page: "policies", - label: "Policies", - icon: , - roles: all_admin_roles, - }, - { - key: "26", - page: "tools", - label: "Tools", - icon: , - children: [ - { key: "18", page: "mcp-servers", label: "MCP Servers", icon: }, - { - key: "21", - page: "vector-stores", - label: "Vector Stores", - icon: , - roles: all_admin_roles, - }, - ], - }, - { - key: "experimental", - page: "experimental", - label: "Experimental", - icon: , - children: [ - { - key: "9", - page: "caching", - label: "Caching", - icon: , - roles: all_admin_roles, - }, - { - key: "25", - page: "prompts", - label: "Prompts", - icon: , - roles: all_admin_roles, - }, - { - key: "10", - page: "budgets", - label: "Budgets", - icon: , - roles: all_admin_roles, - }, - { - key: "20", - page: "transform-request", - label: "API Playground", - icon: , - roles: [...all_admin_roles, ...internalUserRoles], - }, - { - key: "19", - page: "tag-management", - label: "Tag Management", - icon: , - roles: all_admin_roles, - }, - { - key: "27", - page: "claude-code-plugins", - label: "Claude Code Plugins", - icon: , - roles: all_admin_roles, - }, - { key: "4", page: "usage", label: "Old Usage", icon: }, - ], - }, - { - key: "settings", - page: "settings", - label: "Settings", - icon: , - roles: all_admin_roles, - children: [ - { - key: "11", - page: "general-settings", - label: "Router Settings", - icon: , - roles: all_admin_roles, - }, - { - key: "8", - page: "settings", - label: "Logging & Alerts", - icon: , - roles: all_admin_roles, - }, - { - key: "13", - page: "admin-panel", - label: "Admin Settings", - icon: , - roles: all_admin_roles, - }, - { - key: "14", - page: "ui-theme", - label: "UI Theme", - icon: , - roles: all_admin_roles, - }, - ], - }, - ]; + { key: "1", page: "api-keys", label: "Virtual Keys", icon: }, + { + key: "3", + page: "llm-playground", + label: "Test Key", + icon: , + roles: rolesWithWriteAccess, + }, + { + key: "2", + page: "models", + label: "Models + Endpoints", + icon: , + roles: rolesWithWriteAccess, + }, + { + key: "12", + page: "new_usage", + label: "Usage", + icon: , + roles: [...all_admin_roles, ...internalUserRoles], + }, + { key: "6", page: "teams", label: "Teams", icon: }, + { + key: "17", + page: "organizations", + label: "Organizations", + icon: , + roles: all_admin_roles, + }, + { + key: "5", + page: "users", + label: "Internal Users", + icon: , + roles: all_admin_roles, + }, + { key: "14", page: "api_ref", label: "API Reference", icon: }, + { + key: "16", + page: "model-hub-table", + label: "Model Hub", + icon: , + }, + { key: "15", page: "logs", label: "Logs", icon: }, + { + key: "11", + page: "guardrails", + label: "Guardrails", + icon: , + roles: all_admin_roles, + }, + { + key: "28", + page: "policies", + label: "Policies", + icon: , + roles: all_admin_roles, + }, + { + key: "26", + page: "tools", + label: "Tools", + icon: , + children: [ + { key: "18", page: "mcp-servers", label: "MCP Servers", icon: }, + { + key: "21", + page: "vector-stores", + label: "Vector Stores", + icon: , + roles: all_admin_roles, + }, + ], + }, + { + key: "experimental", + page: "experimental", + label: "Experimental", + icon: , + children: [ + { + key: "9", + page: "caching", + label: "Caching", + icon: , + roles: all_admin_roles, + }, + { + key: "25", + page: "prompts", + label: "Prompts", + icon: , + roles: all_admin_roles, + }, + { + key: "10", + page: "budgets", + label: "Budgets", + icon: , + roles: all_admin_roles, + }, + { + key: "20", + page: "transform-request", + label: "API Playground", + icon: , + roles: [...all_admin_roles, ...internalUserRoles], + }, + { + key: "19", + page: "tag-management", + label: "Tag Management", + icon: , + roles: all_admin_roles, + }, + { + key: "27", + page: "claude-code-plugins", + label: "Claude Code Plugins", + icon: , + roles: all_admin_roles, + }, + { key: "4", page: "usage", label: "Old Usage", icon: }, + ], + }, + { + key: "settings", + page: "settings", + label: "Settings", + icon: , + roles: all_admin_roles, + children: [ + { + key: "11", + page: "general-settings", + label: "Router Settings", + icon: , + roles: all_admin_roles, + }, + { + key: "8", + page: "settings", + label: "Logging & Alerts", + icon: , + roles: all_admin_roles, + }, + { + key: "13", + page: "admin-panel", + label: "Admin Settings", + icon: , + roles: all_admin_roles, + }, + { + key: "14", + page: "ui-theme", + label: "UI Theme", + icon: , + roles: all_admin_roles, + }, + ], + }, +]; const Sidebar2: React.FC = ({ accessToken, userRole, defaultSelectedKey, collapsed = false }) => { const router = useRouter(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index 4985206092f..2539cc63f95 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -3,13 +3,14 @@ import { renderHook, waitFor } from "@testing-library/react"; import React, { ReactNode } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { - useModelsInfo, - useModelHub, useAllProxyModels, + useInfiniteModelInfo, + useModelHub, + useModelsInfo, useSelectedTeamModels, - type ProxyModel, type AllProxyModelsResponse, type PaginatedModelInfoResponse, + type ProxyModel, } from "./useModels"; vi.mock("@/components/networking", () => ({ @@ -23,7 +24,7 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => mockUseAuthorized(), })); -import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; +import { modelAvailableCall, modelHubCall, modelInfoCall } from "@/components/networking"; const mockProxyModel: ProxyModel = { id: "model-1", @@ -106,7 +107,7 @@ describe("useModelsInfo", () => { undefined, undefined, undefined, - undefined + undefined, ); expect(modelInfoCall).toHaveBeenCalledTimes(1); }); @@ -130,7 +131,7 @@ describe("useModelsInfo", () => { undefined, undefined, undefined, - undefined + undefined, ); }); @@ -393,7 +394,7 @@ describe("useAllProxyModels", () => { null, true, false, - "expand" + "expand", ); expect(modelAvailableCall).toHaveBeenCalledTimes(1); }); @@ -531,13 +532,7 @@ describe("useSelectedTeamModels", () => { expect(result.current.data).toEqual(mockAllProxyModelsResponse); expect(result.current.error).toBeNull(); - expect(modelAvailableCall).toHaveBeenCalledWith( - "test-access-token", - "test-user-id", - "Admin", - true, - "team-1" - ); + expect(modelAvailableCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", true, "team-1"); expect(modelAvailableCall).toHaveBeenCalledTimes(1); }); @@ -639,3 +634,222 @@ describe("useSelectedTeamModels", () => { expect(modelAvailableCall).not.toHaveBeenCalled(); }); }); + +describe("useInfiniteModelInfo", () => { + let queryClient: QueryClient; + + const mockPageOneResponse: PaginatedModelInfoResponse = { + data: [{ model_name: "gpt-4", model_info: { id: "model-1" } }], + total_count: 2, + current_page: 1, + total_pages: 2, + size: 50, + }; + + const mockPageTwoResponse: PaginatedModelInfoResponse = { + data: [{ model_name: "claude-3", model_info: { id: "model-2" } }], + total_count: 2, + current_page: 2, + total_pages: 2, + size: 50, + }; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return defined result", () => { + (modelInfoCall as any).mockResolvedValue(mockPageOneResponse); + + const { result } = renderHook(() => useInfiniteModelInfo(), { wrapper }); + + expect(result.current).toBeDefined(); + expect(result.current).toHaveProperty("data"); + expect(result.current).toHaveProperty("fetchNextPage"); + expect(result.current).toHaveProperty("hasNextPage"); + expect(result.current).toHaveProperty("isFetchingNextPage"); + expect(result.current).toHaveProperty("isLoading"); + }); + + it("should return paginated data and call modelInfoCall with page 1 initially", async () => { + (modelInfoCall as any).mockResolvedValue(mockPageOneResponse); + + const { result } = renderHook(() => useInfiniteModelInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(true); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data?.pages).toHaveLength(1); + expect(result.current.data?.pages[0]).toEqual(mockPageOneResponse); + expect(result.current.hasNextPage).toBe(true); + expect(modelInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", 1, 50, undefined); + expect(modelInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should use custom size parameter", async () => { + (modelInfoCall as any).mockResolvedValue(mockPageOneResponse); + + const { result } = renderHook(() => useInfiniteModelInfo(25), { wrapper }); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(modelInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", 1, 25, undefined); + }); + + it("should pass search parameter to modelInfoCall", async () => { + (modelInfoCall as any).mockResolvedValue(mockPageOneResponse); + + const { result } = renderHook(() => useInfiniteModelInfo(50, "gpt"), { wrapper }); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(modelInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", 1, 50, "gpt"); + }); + + it("should fetch next page when fetchNextPage is called", async () => { + (modelInfoCall as any).mockResolvedValueOnce(mockPageOneResponse).mockResolvedValueOnce(mockPageTwoResponse); + + const { result } = renderHook(() => useInfiniteModelInfo(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + expect(result.current.hasNextPage).toBe(true); + }); + + await result.current.fetchNextPage(); + + await waitFor(() => { + expect(result.current.data?.pages).toHaveLength(2); + expect(result.current.data?.pages[1]).toEqual(mockPageTwoResponse); + expect(result.current.hasNextPage).toBe(false); + }); + + expect(modelInfoCall).toHaveBeenNthCalledWith(2, "test-access-token", "test-user-id", "Admin", 2, 50, undefined); + }); + + it("should return undefined for hasNextPage when on last page", async () => { + const lastPageResponse: PaginatedModelInfoResponse = { + ...mockPageOneResponse, + current_page: 1, + total_pages: 1, + }; + (modelInfoCall as any).mockResolvedValue(lastPageResponse); + + const { result } = renderHook(() => useInfiniteModelInfo(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.hasNextPage).toBe(false); + }); + + it("should handle error when modelInfoCall fails", async () => { + const errorMessage = "Failed to fetch models"; + const testError = new Error(errorMessage); + (modelInfoCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useInfiniteModelInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(true); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(modelInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useInfiniteModelInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useInfiniteModelInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useInfiniteModelInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelInfoCall).not.toHaveBeenCalled(); + }); +}); 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 c57de675e0e..fe1afdcc39f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -1,4 +1,4 @@ -import { useQuery } from "@tanstack/react-query"; +import { useQuery, useInfiniteQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; import useAuthorized from "../useAuthorized"; @@ -26,6 +26,7 @@ const modelKeys = createQueryKeys("models"); const modelHubKeys = createQueryKeys("modelHub"); const allProxyModelsKeys = createQueryKeys("allProxyModels"); const selectedTeamModelsKeys = createQueryKeys("selectedTeamModels"); +const infiniteModelKeys = createQueryKeys("infiniteModels"); export const useModelsInfo = (page: number = 1, size: number = 50, search?: string, modelId?: string, teamId?: string, sortBy?: string, sortOrder?: string) => { const { accessToken, userId, userRole } = useAuthorized(); @@ -74,3 +75,38 @@ export const useSelectedTeamModels = (teamID: string | null) => { enabled: Boolean(accessToken && userId && userRole && teamID), }); }; + +export const useInfiniteModelInfo = ( + size: number = 50, + search?: string, +) => { + const { accessToken, userId, userRole } = useAuthorized(); + return useInfiniteQuery({ + queryKey: infiniteModelKeys.list({ + filters: { + ...(userId && { userId }), + ...(userRole && { userRole }), + size, + ...(search && { search }), + }, + }), + queryFn: async ({ pageParam }) => { + return await modelInfoCall( + accessToken!, + userId!, + userRole!, + pageParam as number, + size, + search, + ); + }, + initialPageParam: 1, + getNextPageParam: (lastPage) => { + if (lastPage.current_page < lastPage.total_pages) { + return lastPage.current_page + 1; + } + return undefined; + }, + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts index 6429aeafb5a..aba5dddf13d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts @@ -23,6 +23,7 @@ vi.mock("../common/queryKeysFactory", () => ({ // Mock data const mockUIConfig: LiteLLMWellKnownUiConfig = { + sso_configured: true, server_root_path: "/api", proxy_base_url: "https://proxy.example.com", auto_redirect_to_sso: true, @@ -99,6 +100,7 @@ describe("useUIConfig", () => { server_root_path: "/v1", proxy_base_url: null, auto_redirect_to_sso: false, + sso_configured: false, admin_ui_disabled: true, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts index ef4a779b50b..76a3129d6d7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -90,6 +90,7 @@ describe("useAuthorized", () => { proxy_base_url: null, auto_redirect_to_sso: false, admin_ui_disabled: false, + sso_configured: false, }); const decodedPayload = { @@ -131,6 +132,7 @@ describe("useAuthorized", () => { proxy_base_url: null, auto_redirect_to_sso: false, admin_ui_disabled: false, + sso_configured: false, }); decodeTokenMock.mockReturnValue(null); @@ -155,6 +157,7 @@ describe("useAuthorized", () => { proxy_base_url: null, auto_redirect_to_sso: false, admin_ui_disabled: true, + sso_configured: false, }); const decodedPayload = { @@ -190,6 +193,7 @@ describe("useAuthorized", () => { proxy_base_url: null, auto_redirect_to_sso: false, admin_ui_disabled: false, + sso_configured: false, }); decodeTokenMock.mockReturnValue(null); @@ -212,6 +216,7 @@ describe("useAuthorized", () => { proxy_base_url: null, auto_redirect_to_sso: false, admin_ui_disabled: false, + sso_configured: false, }); const decodedPayload = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.test.ts new file mode 100644 index 00000000000..bd0e69c0de3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.test.ts @@ -0,0 +1,190 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import { useDisableUsageIndicator } from "./useDisableUsageIndicator"; +import { LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; + +describe("useDisableUsageIndicator", () => { + const STORAGE_KEY = "disableUsageIndicator"; + + beforeEach(() => { + localStorage.clear(); + vi.clearAllMocks(); + }); + + afterEach(() => { + localStorage.clear(); + }); + + it("should return false when localStorage is empty", () => { + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(false); + }); + + it("should return false when localStorage value is not 'true'", () => { + localStorage.setItem(STORAGE_KEY, "false"); + + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(false); + }); + + it("should return true when localStorage value is 'true'", () => { + localStorage.setItem(STORAGE_KEY, "true"); + + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(true); + }); + + it("should return false when localStorage value is an empty string", () => { + localStorage.setItem(STORAGE_KEY, ""); + + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(false); + }); + + it("should update when storage event fires for the correct key", async () => { + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(false); + + await act(async () => { + localStorage.setItem(STORAGE_KEY, "true"); + const storageEvent = new StorageEvent("storage", { + key: STORAGE_KEY, + newValue: "true", + }); + window.dispatchEvent(storageEvent); + }); + + await waitFor(() => { + expect(result.current).toBe(true); + }); + }); + + it("should not update when storage event fires for a different key", () => { + localStorage.setItem(STORAGE_KEY, "false"); + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(false); + + const storageEvent = new StorageEvent("storage", { + key: "otherKey", + newValue: "true", + }); + window.dispatchEvent(storageEvent); + + expect(result.current).toBe(false); + }); + + it("should update when custom LOCAL_STORAGE_EVENT fires for the correct key", async () => { + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(false); + + await act(async () => { + localStorage.setItem(STORAGE_KEY, "true"); + const customEvent = new CustomEvent(LOCAL_STORAGE_EVENT, { + detail: { key: STORAGE_KEY }, + }); + window.dispatchEvent(customEvent); + }); + + await waitFor(() => { + expect(result.current).toBe(true); + }); + }); + + it("should not update when custom LOCAL_STORAGE_EVENT fires for a different key", () => { + localStorage.setItem(STORAGE_KEY, "false"); + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(false); + + const customEvent = new CustomEvent(LOCAL_STORAGE_EVENT, { + detail: { key: "otherKey" }, + }); + window.dispatchEvent(customEvent); + + expect(result.current).toBe(false); + }); + + it("should update when localStorage changes from false to true via custom event", async () => { + localStorage.setItem(STORAGE_KEY, "false"); + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(false); + + await act(async () => { + localStorage.setItem(STORAGE_KEY, "true"); + const customEvent = new CustomEvent(LOCAL_STORAGE_EVENT, { + detail: { key: STORAGE_KEY }, + }); + window.dispatchEvent(customEvent); + }); + + await waitFor(() => { + expect(result.current).toBe(true); + }); + }); + + it("should update when localStorage changes from true to false via storage event", async () => { + localStorage.setItem(STORAGE_KEY, "true"); + const { result } = renderHook(() => useDisableUsageIndicator()); + + expect(result.current).toBe(true); + + await act(async () => { + localStorage.setItem(STORAGE_KEY, "false"); + const storageEvent = new StorageEvent("storage", { + key: STORAGE_KEY, + newValue: "false", + }); + window.dispatchEvent(storageEvent); + }); + + await waitFor(() => { + expect(result.current).toBe(false); + }); + }); + + it("should cleanup event listeners on unmount", () => { + const addEventListenerSpy = vi.spyOn(window, "addEventListener"); + const removeEventListenerSpy = vi.spyOn(window, "removeEventListener"); + + const { unmount } = renderHook(() => useDisableUsageIndicator()); + + expect(addEventListenerSpy).toHaveBeenCalledTimes(2); + expect(addEventListenerSpy).toHaveBeenCalledWith("storage", expect.any(Function)); + expect(addEventListenerSpy).toHaveBeenCalledWith(LOCAL_STORAGE_EVENT, expect.any(Function)); + + unmount(); + + expect(removeEventListenerSpy).toHaveBeenCalledTimes(2); + expect(removeEventListenerSpy).toHaveBeenCalledWith("storage", expect.any(Function)); + expect(removeEventListenerSpy).toHaveBeenCalledWith(LOCAL_STORAGE_EVENT, expect.any(Function)); + }); + + it("should handle multiple hooks independently", async () => { + const { result: result1 } = renderHook(() => useDisableUsageIndicator()); + const { result: result2 } = renderHook(() => useDisableUsageIndicator()); + + expect(result1.current).toBe(false); + expect(result2.current).toBe(false); + + await act(async () => { + localStorage.setItem(STORAGE_KEY, "true"); + const customEvent = new CustomEvent(LOCAL_STORAGE_EVENT, { + detail: { key: STORAGE_KEY }, + }); + window.dispatchEvent(customEvent); + }); + + await waitFor(() => { + expect(result1.current).toBe(true); + expect(result2.current).toBe(true); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.ts new file mode 100644 index 00000000000..7f4e2295090 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableUsageIndicator.ts @@ -0,0 +1,33 @@ +import { getLocalStorageItem, LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; +import { useSyncExternalStore } from "react"; + +function subscribe(callback: () => void) { + const onStorage = (e: StorageEvent) => { + if (e.key === "disableUsageIndicator") { + callback(); + } + }; + + const onCustom = (e: Event) => { + const { key } = (e as CustomEvent).detail; + if (key === "disableUsageIndicator") { + callback(); + } + }; + + window.addEventListener("storage", onStorage); + window.addEventListener(LOCAL_STORAGE_EVENT, onCustom); + + return () => { + window.removeEventListener("storage", onStorage); + window.removeEventListener(LOCAL_STORAGE_EVENT, onCustom); + }; +} + +function getSnapshot() { + return getLocalStorageItem("disableUsageIndicator") === "true"; +} + +export function useDisableUsageIndicator() { + return useSyncExternalStore(subscribe, getSnapshot); +} diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx index 79834512605..ad2dde2da83 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx @@ -64,7 +64,12 @@ describe("LoginPage", () => { it("should render", async () => { (useUIConfig as ReturnType).mockReturnValue({ - data: { auto_redirect_to_sso: false, server_root_path: "/", proxy_base_url: null }, + data: { + auto_redirect_to_sso: false, + server_root_path: "/", + proxy_base_url: null, + sso_configured: false, + }, isLoading: false, }); (getCookie as ReturnType).mockReturnValue(null); @@ -84,7 +89,12 @@ describe("LoginPage", () => { it("should call router.replace to dashboard when jwt is valid", async () => { const validToken = "valid-token"; (useUIConfig as ReturnType).mockReturnValue({ - data: { auto_redirect_to_sso: false, server_root_path: "/", proxy_base_url: null }, + data: { + auto_redirect_to_sso: false, + server_root_path: "/", + proxy_base_url: null, + sso_configured: false, + }, isLoading: false, }); (getCookie as ReturnType).mockReturnValue(validToken); @@ -105,7 +115,12 @@ describe("LoginPage", () => { it("should call router.push to SSO when jwt is invalid and auto_redirect_to_sso is true", async () => { const invalidToken = "invalid-token"; (useUIConfig as ReturnType).mockReturnValue({ - data: { auto_redirect_to_sso: true, server_root_path: "/", proxy_base_url: null }, + data: { + auto_redirect_to_sso: true, + server_root_path: "/", + proxy_base_url: null, + sso_configured: true, + }, isLoading: false, }); (getCookie as ReturnType).mockReturnValue(invalidToken); @@ -126,7 +141,12 @@ describe("LoginPage", () => { it("should not call router when jwt is invalid and auto_redirect_to_sso is false", async () => { const invalidToken = "invalid-token"; (useUIConfig as ReturnType).mockReturnValue({ - data: { auto_redirect_to_sso: false, server_root_path: "/", proxy_base_url: null }, + data: { + auto_redirect_to_sso: false, + server_root_path: "/", + proxy_base_url: null, + sso_configured: false, + }, isLoading: false, }); (getCookie as ReturnType).mockReturnValue(invalidToken); @@ -150,7 +170,12 @@ describe("LoginPage", () => { it("should send user to dashboard when jwt is valid even if auto_redirect_to_sso is true", async () => { const validToken = "valid-token"; (useUIConfig as ReturnType).mockReturnValue({ - data: { auto_redirect_to_sso: true, server_root_path: "/", proxy_base_url: null }, + data: { + auto_redirect_to_sso: true, + server_root_path: "/", + proxy_base_url: null, + sso_configured: true, + }, isLoading: false, }); (getCookie as ReturnType).mockReturnValue(validToken); @@ -172,7 +197,12 @@ describe("LoginPage", () => { it("should show alert when admin_ui_disabled is true", async () => { (useUIConfig as ReturnType).mockReturnValue({ - data: { admin_ui_disabled: true, server_root_path: "/", proxy_base_url: null }, + data: { + admin_ui_disabled: true, + server_root_path: "/", + proxy_base_url: null, + sso_configured: false, + }, isLoading: false, }); (getCookie as ReturnType).mockReturnValue(null); @@ -192,4 +222,60 @@ describe("LoginPage", () => { expect(mockPush).not.toHaveBeenCalled(); expect(mockReplace).not.toHaveBeenCalled(); }); + + it("should show Login with SSO button when sso_configured is true", async () => { + (useUIConfig as ReturnType).mockReturnValue({ + data: { + auto_redirect_to_sso: false, + server_root_path: "/", + proxy_base_url: null, + sso_configured: true, + }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(null); + (isJwtExpired as ReturnType).mockReturnValue(true); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(screen.getByRole("heading", { name: "Login" })).toBeInTheDocument(); + }); + + expect(screen.getByRole("button", { name: "Login with SSO" })).toBeInTheDocument(); + }); + + it("should show disabled Login with SSO button with popover when sso_configured is false", async () => { + (useUIConfig as ReturnType).mockReturnValue({ + data: { + auto_redirect_to_sso: false, + server_root_path: "/", + proxy_base_url: null, + sso_configured: false, + }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(null); + (isJwtExpired as ReturnType).mockReturnValue(true); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(screen.getByRole("heading", { name: "Login" })).toBeInTheDocument(); + }); + + const ssoButton = screen.getByRole("button", { name: "Login with SSO" }); + expect(ssoButton).toBeInTheDocument(); + expect(ssoButton).toBeDisabled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index 620cb41dfee..a05fa4e214e 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -8,7 +8,7 @@ import { getCookie } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; import { InfoCircleOutlined } from "@ant-design/icons"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { Alert, Button, Card, Form, Input, Space, Typography } from "antd"; +import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd"; import { useRouter } from "next/navigation"; import { useEffect, useState } from "react"; @@ -179,8 +179,39 @@ function LoginPageContent() { {isLoginLoading ? "Logging in..." : "Login"} + + {!uiConfig?.sso_configured ? ( + + + + ) : ( + + )} + + {uiConfig?.sso_configured && ( + Single Sign-On (SSO) is enabled. LiteLLM no longer automatically redirects to the SSO login flow upon loading this page. To re-enable auto-redirect-to-SSO, set AUTO_REDIRECT_UI_LOGIN_TO_SSO=true in your environment configuration.} + /> + )} ); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx index 0a5cd17e571..ee59ac84ece 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx @@ -71,6 +71,7 @@ describe("ModelHubTable", () => { proxy_base_url: "http://localhost:4000", auto_redirect_to_sso: false, admin_ui_disabled: false, + sso_configured: false, }); vi.mocked(networking.modelHubPublicModelsCall).mockResolvedValue([]); vi.mocked(networking.getUiSettings).mockResolvedValue({ @@ -140,6 +141,7 @@ describe("ModelHubTable", () => { proxy_base_url: "http://localhost:4000", auto_redirect_to_sso: false, admin_ui_disabled: false, + sso_configured: false, }); modelHubPublicModelsCallMock.mockResolvedValue([]); vi.mocked(networking.getUiSettings).mockResolvedValue({ diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx index 2b7399c4565..74b2619f7f3 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx @@ -200,6 +200,7 @@ export const ModelSelect = (props: ModelSelectProps) => { }, ]} mode="multiple" + placeholder="Select Models" allowClear maxTagCount="responsive" maxTagPlaceholder={(omittedValues) => ( diff --git a/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.test.tsx b/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.test.tsx new file mode 100644 index 00000000000..b91f97885e6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.test.tsx @@ -0,0 +1,301 @@ +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; +import { PaginatedModelSelect } from "./PaginatedModelSelect"; + +const mockFetchNextPage = vi.fn(); + +vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ + useInfiniteModelInfo: vi.fn(), +})); + +vi.mock("@tanstack/react-pacer/debouncer", () => { + const React = require("react"); + return { + useDebouncedState: (initial: string) => { + const [value, setValue] = React.useState(initial); + return [value, setValue]; + }, + }; +}); + +import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; + +const mockUseInfiniteModelInfo = vi.mocked(useInfiniteModelInfo); + +const mockPagesWithModels = { + pages: [ + { + data: [ + { model_name: "GPT-4", model_info: { id: "model-1" } }, + { model_name: "Claude-3", model_info: { id: "model-2" } }, + ], + total_count: 2, + current_page: 1, + total_pages: 1, + size: 50, + }, + ], +}; + +const mockEmptyPages = { + pages: [{ data: [], total_count: 0, current_page: 1, total_pages: 1, size: 50 }], +}; + +describe("PaginatedModelSelect", () => { + const mockOnChange = vi.fn(); + + const defaultHookReturn = { + data: mockPagesWithModels, + fetchNextPage: mockFetchNextPage, + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockUseInfiniteModelInfo.mockReturnValue(defaultHookReturn as any); + }); + + it("should render", () => { + renderWithProviders(); + + expect(screen.getByRole("combobox")).toBeInTheDocument(); + expect(screen.getByText("Select a model")).toBeInTheDocument(); + }); + + it("should display custom placeholder when provided", () => { + renderWithProviders( + , + ); + + expect(screen.getByText("Choose model")).toBeInTheDocument(); + }); + + it("should display model options when data is loaded", async () => { + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + await userEvent.click(combobox); + + await waitFor(() => { + expect(screen.getByRole("option", { name: "GPT-4 (model-1)" })).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "Claude-3 (model-2)" })).toBeInTheDocument(); + }); + }); + + it("should call onChange when user selects a model", async () => { + const user = userEvent.setup({ delay: null }); + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + await user.click(combobox); + + const visibleOption = await screen.findByTitle("GPT-4 (model-1)"); + await user.click(visibleOption); + + await waitFor(() => { + expect(mockOnChange).toHaveBeenCalledWith("model-1"); + }); + }); + + it("should display selected value when value prop is provided", async () => { + renderWithProviders( + , + ); + + const combobox = screen.getByRole("combobox"); + await userEvent.click(combobox); + + await waitFor(() => { + expect(screen.getByRole("option", { name: "GPT-4 (model-1)" })).toBeInTheDocument(); + }); + }); + + it("should show loading state when isLoading is true", () => { + mockUseInfiniteModelInfo.mockReturnValue({ + ...defaultHookReturn, + isLoading: true, + } as any); + + renderWithProviders(); + + expect(screen.getByRole("combobox")).toHaveAttribute("aria-expanded", "false"); + }); + + it("should pass pageSize to useInfiniteModelInfo", () => { + renderWithProviders(); + + expect(mockUseInfiniteModelInfo).toHaveBeenCalledWith(25, undefined); + }); + + it("should pass search to useInfiniteModelInfo when user types", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + await user.click(combobox); + await user.keyboard("gpt"); + + await waitFor(() => { + expect(mockUseInfiniteModelInfo).toHaveBeenCalledWith(50, "gpt"); + }); + }); + + it("should have scroll container for infinite loading when hasNextPage is true", async () => { + mockUseInfiniteModelInfo.mockReturnValue({ + ...defaultHookReturn, + hasNextPage: true, + isFetchingNextPage: false, + } as any); + + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + await userEvent.click(combobox); + + await waitFor(() => { + expect(screen.getByRole("option", { name: "GPT-4 (model-1)" })).toBeInTheDocument(); + }); + + const scrollableContainer = document.querySelector( + ".ant-select-dropdown .rc-virtual-list-holder", + ); + expect(scrollableContainer).toBeInTheDocument(); + expect(scrollableContainer).toHaveAttribute("style"); + }); + + it("should deduplicate models with same id across pages", async () => { + mockUseInfiniteModelInfo.mockReturnValue({ + ...defaultHookReturn, + data: { + pages: [ + { + data: [ + { model_name: "GPT-4", model_info: { id: "model-1" } }, + { model_name: "GPT-4 Dupe", model_info: { id: "model-1" } }, + ], + total_count: 2, + current_page: 1, + total_pages: 1, + size: 50, + }, + ], + }, + fetchNextPage: mockFetchNextPage, + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + } as any); + + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + await userEvent.click(combobox); + + await waitFor(() => { + const model1Options = screen.queryAllByRole("option", { name: /model-1/ }); + expect(model1Options.length).toBe(1); + }); + }); + + it("should skip models without model_info id", async () => { + mockUseInfiniteModelInfo.mockReturnValue({ + ...defaultHookReturn, + data: { + pages: [ + { + data: [ + { model_name: "Valid Model", model_info: { id: "valid-id" } }, + { model_name: "No ID", model_info: null }, + { model_name: "Empty ID", model_info: { id: "" } }, + ], + total_count: 3, + current_page: 1, + total_pages: 1, + size: 50, + }, + ], + }, + fetchNextPage: mockFetchNextPage, + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + } as any); + + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + await userEvent.click(combobox); + + await waitFor(() => { + expect(screen.getByRole("option", { name: "Valid Model (valid-id)" })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: "No ID" })).not.toBeInTheDocument(); + expect(screen.queryByRole("option", { name: "Empty ID" })).not.toBeInTheDocument(); + }); + }); + + it("should show model ID only when model_name is empty", async () => { + mockUseInfiniteModelInfo.mockReturnValue({ + ...defaultHookReturn, + data: { + pages: [ + { + data: [{ model_name: "", model_info: { id: "id-only" } }], + total_count: 1, + current_page: 1, + total_pages: 1, + size: 50, + }, + ], + }, + fetchNextPage: mockFetchNextPage, + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + } as any); + + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + await userEvent.click(combobox); + + await waitFor(() => { + expect(screen.getByRole("option", { name: "id-only" })).toBeInTheDocument(); + }); + }); + + it("should respect allowClear prop", () => { + renderWithProviders( + , + ); + + expect(screen.getByRole("combobox")).toBeInTheDocument(); + }); + + it("should respect disabled prop", () => { + renderWithProviders(); + + const combobox = screen.getByRole("combobox"); + expect(combobox.closest(".ant-select")).toHaveClass("ant-select-disabled"); + }); + + it("should not call fetchNextPage when hasNextPage is false", async () => { + mockUseInfiniteModelInfo.mockReturnValue({ + ...defaultHookReturn, + hasNextPage: false, + } as any); + + renderWithProviders(); + + await userEvent.click(screen.getByRole("combobox")); + + await waitFor(() => { + expect(screen.getByRole("option", { name: "GPT-4 (model-1)" })).toBeInTheDocument(); + }); + + expect(mockFetchNextPage).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx new file mode 100644 index 00000000000..9b22fd1bd87 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelSelect/PaginatedModelSelect/PaginatedModelSelect.tsx @@ -0,0 +1,143 @@ +import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; +import { LoadingOutlined } from "@ant-design/icons"; +import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; +import { Select, Space, Typography } from "antd"; +import { useMemo, useState, type UIEvent } from "react"; + +const { Text } = Typography; + +export interface PaginatedModelSelectProps { + value?: string; + onChange?: (value: string) => void; + placeholder?: string; + style?: React.CSSProperties; + pageSize?: number; + allowClear?: boolean; + disabled?: boolean; +} + +const SCROLL_THRESHOLD = 0.8; +const DEBOUNCE_MS = 300; + +export const PaginatedModelSelect = ({ + value, + onChange, + placeholder = "Select a model", + style, + pageSize = 50, + allowClear = true, + disabled = false, +}: PaginatedModelSelectProps) => { + const [searchInput, setSearchInput] = useState(""); + const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", { + wait: DEBOUNCE_MS, + }); + + const { + data, + fetchNextPage, + hasNextPage, + isFetchingNextPage, + isLoading, + } = useInfiniteModelInfo(pageSize, debouncedSearch || undefined); + + const options = useMemo(() => { + if (!data?.pages) return []; + + const seen = new Set(); + const result: { label: string; value: string; modelName: string; modelId: string }[] = []; + + for (const page of data.pages) { + for (const model of page.data) { + const modelId = model.model_info?.id ?? ""; + const modelName = model.model_name ?? ""; + + // Dedupe by id - skip models without id (can't uniquely identify) + if (!modelId || seen.has(modelId)) continue; + seen.add(modelId); + + result.push({ + label: modelName ? `${modelName} (${modelId})` : modelId, + value: modelId, + modelName, + modelId, + }); + } + } + + return result; + }, [data]); + + const optionRender = (option: { data: { modelName: string; modelId: string; label: string } }) => { + const { modelName, modelId } = option.data; + + return ( + <> + {modelName ? ( + + + Model name: + {modelName} + + + Model ID: {modelId} + + + ) : ( + Model ID: {modelId} + )} + + ); + }; + + const handlePopupScroll = (e: UIEvent) => { + const target = e.currentTarget; + const scrollRatio = + (target.scrollTop + target.clientHeight) / target.scrollHeight; + + if (scrollRatio >= SCROLL_THRESHOLD && hasNextPage && !isFetchingNextPage) { + fetchNextPage(); + } + }; + + const handleSearch = (value: string) => { + setSearchInput(value); + setDebouncedSearch(value); + }; + + const handleChange = (v: string | string[] | null) => { + const normalized = + typeof v === "string" ? v : Array.isArray(v) ? v[0] ?? "" : ""; + onChange?.(normalized); + }; + + return ( + = ({ visible, onClose, a )) || ( - <> - - - - - - )} + + + )} @@ -721,21 +720,21 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a const renderStepButtons = () => { const totalSteps = shouldRenderContentFilterConfigSettings(selectedProvider) ? 4 : 2; const isLastStep = currentStep === totalSteps - 1; - + return (
{currentStep > 0 && ( - )} - {!isLastStep && } + {!isLastStep && } {isLastStep && ( - )} -
diff --git a/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentCategoryConfiguration.tsx b/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentCategoryConfiguration.tsx index 7c8652a7b18..3b0b0acce85 100644 --- a/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentCategoryConfiguration.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentCategoryConfiguration.tsx @@ -1,12 +1,10 @@ import React from "react"; -import { Card, Typography, Select, Table, Tag, Collapse } from "antd"; +import { Card, Typography, Select, Table, Tag, Collapse, Button } from "antd"; import { DeleteOutlined, PlusOutlined, FileTextOutlined } from "@ant-design/icons"; -import { Button } from "@tremor/react"; import { getCategoryYaml } from "../../networking"; const { Title, Text } = Typography; const { Option } = Select; -const { Panel } = Collapse; interface ContentCategory { name: string; @@ -191,10 +189,9 @@ const ContentCategoryConfiguration: React.FC width: 80, render: (_: any, record: SelectedCategory) => ( @@ -244,9 +241,10 @@ const ContentCategoryConfiguration: React.FC ))} @@ -308,7 +306,6 @@ const ContentCategoryConfiguration: React.FC activeKey={expandedYamlCategories} onChange={(keys) => { const keyArray = Array.isArray(keys) ? keys : keys ? [keys] : []; - const newExpanded = new Set(keyArray as string[]); const oldExpanded = new Set(expandedYamlCategories); // Find newly expanded categories and fetch their YAML @@ -322,44 +319,40 @@ const ContentCategoryConfiguration: React.FC setExpandedYamlCategories(keyArray as string[]); }} ghost - > - {selectedCategories.map((category) => ( - - - View YAML for {category.display_name} - - } - key={category.category} - > - {loadingYaml[category.category] ? ( -
- Loading YAML... -
- ) : categoryYaml[category.category] ? ( - 
-                      {categoryYaml[category.category]}
-
- ) : ( -
- YAML will load when expanded -
- )} - - ))} - + items={selectedCategories.map((category) => ({ + key: category.category, + label: ( +
+ + View YAML for {category.display_name} +
+ ), + children: loadingYaml[category.category] ? ( +
+ Loading YAML... +
+ ) : categoryYaml[category.category] ? ( + 
+                    {categoryYaml[category.category]}
+
+ ) : ( +
+ YAML will load when expanded +
+ ), + }))} + /> ) : ( diff --git a/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterConfiguration.tsx b/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterConfiguration.tsx index bae95aac6ce..882abc0b933 100644 --- a/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterConfiguration.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterConfiguration.tsx @@ -1,7 +1,6 @@ import React, { useState } from "react"; -import { Typography, Space, Upload, Card } from "antd"; +import { Typography, Space, Upload, Card, Button } from "antd"; import { PlusOutlined, UploadOutlined } from "@ant-design/icons"; -import { Button } from "@tremor/react"; import { validateBlockedWordsFile } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; import PatternModal from "./PatternModal"; @@ -221,10 +220,10 @@ const ContentFilterConfiguration: React.FC = ({ >
- - @@ -253,11 +252,11 @@ const ContentFilterConfiguration: React.FC = ({ >
- - diff --git a/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterManager.test.tsx b/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterManager.test.tsx index 6d879e1c544..cbd7033237c 100644 --- a/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterManager.test.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterManager.test.tsx @@ -1,58 +1,195 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { render, screen, waitFor } from "@testing-library/react"; -import ContentFilterManager from "./ContentFilterManager"; +import userEvent from "@testing-library/user-event"; +import ContentFilterManager, { + formatContentFilterDataForAPI, +} from "./ContentFilterManager"; import React from "react"; +const CONTENT_FILTER_GUARDRAIL_DATA = { + litellm_params: { + guardrail: "litellm_content_filter", + patterns: [{ pattern_type: "prebuilt", pattern_name: "email", action: "BLOCK" }], + blocked_words: [{ keyword: "test", action: "BLOCK", description: null }], + }, +}; + +const GUARDRAIL_SETTINGS = { + content_filter_settings: { + prebuilt_patterns: [], + pattern_categories: ["PII"], + supported_actions: ["BLOCK", "MASK"], + }, +}; + vi.mock("./ContentFilterConfiguration", () => ({ - default: () =>
Mock Content Filter Configuration
+ default: ({ + onPatternAdd, + onPatternRemove, + onBlockedWordAdd, + onBlockedWordRemove, + selectedPatterns, + blockedWords, + }: { + onPatternAdd: (p: object) => void; + onPatternRemove: (id: string) => void; + onBlockedWordAdd: (w: object) => void; + onBlockedWordRemove: (id: string) => void; + selectedPatterns: { id: string }[]; + blockedWords: { id: string }[]; + }) => ( +
+ + + {selectedPatterns[0] && ( + + )} + {blockedWords[0] && ( + + )} +
+ ), })); vi.mock("./ContentFilterDisplay", () => ({ - default: () =>
Mock Content Filter Display
+ default: ({ + patterns, + blockedWords, + }: { + patterns: { name: string }[]; + blockedWords: { keyword: string }[]; + }) => ( +
+ Patterns: {patterns.map((p) => p.name).join(", ")} + Keywords: {blockedWords.map((w) => w.keyword).join(", ")} +
+ ), })); -vi.mock("antd", () => ({ - Divider: ({ children }: { children: React.ReactNode }) =>
{children}
-})); +vi.mock("antd", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + Divider: ({ children }: { children: React.ReactNode }) => ( +
{children}
+ ), + Alert: ({ + message, + type, + }: { + message: React.ReactNode; + type: string; + }) => ( +
+ {message} +
+ ), + }; +}); -describe("ContentFilterManager - Unsaved Changes Detection", () => { +describe("ContentFilterManager", () => { beforeEach(() => { vi.clearAllMocks(); }); + it("should render when guardrail is content filter and isEditing is true", async () => { + render( + + ); + + await waitFor(() => { + expect(screen.getByTestId("content-filter-config")).toBeInTheDocument(); + }); + + expect(screen.getByTestId("divider")).toHaveTextContent( + "Content Filter Configuration" + ); + }); + + it("should return null when guardrail is not litellm_content_filter", () => { + const guardrailData = { + litellm_params: { guardrail: "presidio" }, + }; + + const { container } = render( + + ); + + expect(screen.queryByTestId("content-filter-config")).not.toBeInTheDocument(); + expect(screen.queryByTestId("content-filter-display")).not.toBeInTheDocument(); + expect(container.firstChild).toBeNull(); + }); + + it("should render read-only display when isEditing is false", async () => { + render( + + ); + + await waitFor(() => { + expect(screen.getByTestId("content-filter-display")).toBeInTheDocument(); + }); + + expect(screen.queryByTestId("content-filter-config")).not.toBeInTheDocument(); + expect(screen.getByText(/email/)).toBeInTheDocument(); + expect(screen.getByText(/test/)).toBeInTheDocument(); + }); + it("should call onUnsavedChanges with false when component initializes with matching data", async () => { - /** - * Tests that the ContentFilterManager correctly initializes the unsaved changes - * detection and calls onUnsavedChanges(false) when the current state matches - * the original loaded state (no changes yet). - */ const mockOnUnsavedChanges = vi.fn(); const mockOnDataChange = vi.fn(); - const guardrailData = { - litellm_params: { - guardrail: "litellm_content_filter", - patterns: [ - { pattern_type: "prebuilt", pattern_name: "email", action: "BLOCK" } - ], - blocked_words: [ - { keyword: "test", action: "BLOCK", description: null } - ] - } - }; - - const guardrailSettings = { - content_filter_settings: { - prebuilt_patterns: [], - pattern_categories: ["PII"], - supported_actions: ["BLOCK", "MASK"] - } - }; - render( { /> ); - // Wait for component to render in edit mode await waitFor(() => { expect(screen.getByTestId("content-filter-config")).toBeInTheDocument(); }); - // Verify onUnsavedChanges was called with false (no changes initially) await waitFor(() => { expect(mockOnUnsavedChanges).toHaveBeenCalledWith(false); }); - // Verify onDataChange was called with initial data expect(mockOnDataChange).toHaveBeenCalled(); }); + + it("should call onUnsavedChanges with true when user adds a pattern", async () => { + const mockOnUnsavedChanges = vi.fn(); + const user = userEvent.setup(); + + render( + + ); + + await waitFor(() => { + expect(screen.getByTestId("content-filter-config")).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /add pattern/i })); + + await waitFor(() => { + expect(mockOnUnsavedChanges).toHaveBeenCalledWith(true); + }); + }); + + it("should call onUnsavedChanges with true when user adds a keyword", async () => { + const mockOnUnsavedChanges = vi.fn(); + const user = userEvent.setup(); + + render( + + ); + + await waitFor(() => { + expect(screen.getByTestId("content-filter-config")).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /add keyword/i })); + + await waitFor(() => { + expect(mockOnUnsavedChanges).toHaveBeenCalledWith(true); + }); + }); + + it("should call onUnsavedChanges with true when user removes a pattern", async () => { + const mockOnUnsavedChanges = vi.fn(); + const user = userEvent.setup(); + + render( + + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /remove pattern/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /remove pattern/i })); + + await waitFor(() => { + expect(mockOnUnsavedChanges).toHaveBeenCalledWith(true); + }); + }); + + it("should show unsaved changes alert when data has changed", async () => { + const user = userEvent.setup(); + + render( + + ); + + await waitFor(() => { + expect(screen.getByTestId("content-filter-config")).toBeInTheDocument(); + }); + + expect(screen.queryByTestId("unsaved-alert")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: /add pattern/i })); + + await waitFor(() => { + expect(screen.getByTestId("unsaved-alert")).toBeInTheDocument(); + }); + + expect(screen.getByTestId("unsaved-alert")).toHaveTextContent( + /unsaved changes.*Save Changes/i + ); + }); + + it("should call onDataChange when patterns or keywords change", async () => { + const mockOnDataChange = vi.fn(); + const user = userEvent.setup(); + + render( + + ); + + await waitFor(() => { + expect(mockOnDataChange).toHaveBeenCalled(); + }); + + const initialCalls = mockOnDataChange.mock.calls.length; + await user.click(screen.getByRole("button", { name: /add keyword/i })); + + await waitFor(() => { + expect(mockOnDataChange.mock.calls.length).toBeGreaterThan(initialCalls); + }); + + const lastCall = mockOnDataChange.mock.calls[mockOnDataChange.mock.calls.length - 1]; + const blockedWords = lastCall[1]; + expect(blockedWords).toContainEqual( + expect.objectContaining({ keyword: "secret", action: "MASK" }) + ); + }); + + it("should not call onUnsavedChanges when isEditing is false", async () => { + const mockOnUnsavedChanges = vi.fn(); + + render( + + ); + + await waitFor(() => { + expect(screen.getByTestId("content-filter-display")).toBeInTheDocument(); + }); + + expect(mockOnUnsavedChanges).not.toHaveBeenCalled(); + }); + + it("should initialize with empty data when guardrailData has no patterns or blocked_words", async () => { + const guardrailData = { + litellm_params: { + guardrail: "litellm_content_filter", + }, + }; + + const mockOnDataChange = vi.fn(); + + render( + + ); + + await waitFor(() => { + expect(mockOnDataChange).toHaveBeenCalledWith([], [], []); + }); + }); + + it("should not render ContentFilterConfiguration when guardrailSettings has no content_filter_settings", async () => { + render( + + ); + + await waitFor(() => { + expect(screen.getByTestId("divider")).toBeInTheDocument(); + }); + + expect(screen.queryByTestId("content-filter-config")).not.toBeInTheDocument(); + }); }); +describe("formatContentFilterDataForAPI", () => { + it("should format patterns and blocked words for API", () => { + const patterns = [ + { + id: "p1", + type: "prebuilt" as const, + name: "email", + action: "BLOCK" as const, + }, + { + id: "p2", + type: "custom" as const, + name: "custom", + pattern: "\\d+", + action: "MASK" as const, + }, + ]; + const blockedWords = [ + { + id: "w1", + keyword: "secret", + action: "MASK" as const, + description: "-sensitive", + }, + ]; + + const result = formatContentFilterDataForAPI(patterns, blockedWords); + + expect(result.patterns).toEqual([ + { + pattern_type: "prebuilt", + pattern_name: "email", + pattern: undefined, + name: "email", + action: "BLOCK", + }, + { + pattern_type: "regex", + pattern_name: undefined, + pattern: "\\d+", + name: "custom", + action: "MASK", + }, + ]); + expect(result.blocked_words).toEqual([ + { keyword: "secret", action: "MASK", description: "-sensitive" }, + ]); + expect(result.categories).toBeUndefined(); + }); + + it("should include categories when provided", () => { + const patterns: Parameters[0] = []; + const blockedWords: Parameters[1] = []; + const categories = [ + { + id: "c1", + category: "PII", + display_name: "PII", + action: "BLOCK" as const, + severity_threshold: "high" as const, + }, + ]; + + const result = formatContentFilterDataForAPI(patterns, blockedWords, categories); + + expect(result.categories).toEqual([ + { + category: "PII", + enabled: true, + action: "BLOCK", + severity_threshold: "high", + }, + ]); + }); + + it("should use medium as default severity_threshold when category has none", () => { + const patterns: Parameters[0] = []; + const blockedWords: Parameters[1] = []; + const categories = [ + { + id: "c1", + category: "PII", + display_name: "PII", + action: "MASK" as const, + severity_threshold: undefined as unknown as "high" | "medium" | "low", + }, + ]; + + const result = formatContentFilterDataForAPI(patterns, blockedWords, categories); + + expect(result.categories).toEqual([ + { + category: "PII", + enabled: true, + action: "MASK", + severity_threshold: "medium", + }, + ]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterManager.tsx b/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterManager.tsx index 3fac7e11393..aa23c0e1db0 100644 --- a/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterManager.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/content_filter/ContentFilterManager.tsx @@ -1,8 +1,10 @@ -import React, { useState, useEffect } from "react"; -import { Divider } from "antd"; +import { Alert, Divider, Typography } from "antd"; +import React, { useEffect, useState } from "react"; import ContentFilterConfiguration from "./ContentFilterConfiguration"; import ContentFilterDisplay from "./ContentFilterDisplay"; +const { Text } = Typography + interface Pattern { id: string; type: "prebuilt" | "custom"; @@ -19,6 +21,21 @@ interface BlockedWord { description?: string; } +interface SelectedContentCategory { + id: string; + category: string; + display_name: string; + action: "BLOCK" | "MASK"; + severity_threshold: "high" | "medium" | "low"; +} + +interface ContentCategory { + name: string; + display_name: string; + description: string; + default_action: string; +} + interface GuardrailSettings { content_filter_settings?: { prebuilt_patterns: Array<{ @@ -29,6 +46,7 @@ interface GuardrailSettings { }>; pattern_categories: string[]; supported_actions: string[]; + content_categories?: ContentCategory[]; }; } @@ -37,7 +55,7 @@ interface ContentFilterManagerProps { guardrailSettings: GuardrailSettings | null; isEditing: boolean; accessToken: string | null; - onDataChange?: (patterns: Pattern[], blockedWords: BlockedWord[]) => void; + onDataChange?: (patterns: Pattern[], blockedWords: BlockedWord[], categories: SelectedContentCategory[]) => void; onUnsavedChanges?: (hasChanges: boolean) => void; } @@ -51,8 +69,10 @@ const ContentFilterManager: React.FC = ({ }) => { const [selectedPatterns, setSelectedPatterns] = useState([]); const [blockedWords, setBlockedWords] = useState([]); + const [selectedContentCategories, setSelectedContentCategories] = useState([]); const [originalPatterns, setOriginalPatterns] = useState([]); const [originalBlockedWords, setOriginalBlockedWords] = useState([]); + const [originalContentCategories, setOriginalContentCategories] = useState([]); // Load data from guardrail on mount or when guardrailData changes useEffect(() => { @@ -85,21 +105,45 @@ const ContentFilterManager: React.FC = ({ setBlockedWords([]); setOriginalBlockedWords([]); } - }, [guardrailData]); + + if (guardrailData?.litellm_params?.categories?.length > 0) { + const contentCategoriesMap = guardrailSettings?.content_filter_settings?.content_categories + ? Object.fromEntries( + guardrailSettings.content_filter_settings.content_categories.map((c) => [c.name, c]) + ) + : {}; + const categories = guardrailData.litellm_params.categories.map((c: any, index: number) => { + const meta = contentCategoriesMap[c.category]; + return { + id: `category-${index}`, + category: c.category, + display_name: meta?.display_name ?? c.category, + action: (c.action || "BLOCK") as "BLOCK" | "MASK", + severity_threshold: (c.severity_threshold || "medium") as "high" | "medium" | "low", + }; + }); + setSelectedContentCategories(categories); + setOriginalContentCategories(categories); + } else { + setSelectedContentCategories([]); + setOriginalContentCategories([]); + } + }, [guardrailData, guardrailSettings?.content_filter_settings?.content_categories]); // Notify parent component when data changes useEffect(() => { if (onDataChange) { - onDataChange(selectedPatterns, blockedWords); + onDataChange(selectedPatterns, blockedWords, selectedContentCategories); } - }, [selectedPatterns, blockedWords, onDataChange]); + }, [selectedPatterns, blockedWords, selectedContentCategories, onDataChange]); // Detect unsaved changes const hasUnsavedChanges = React.useMemo(() => { const hasPatternChanges = JSON.stringify(selectedPatterns) !== JSON.stringify(originalPatterns); const hasWordChanges = JSON.stringify(blockedWords) !== JSON.stringify(originalBlockedWords); - return hasPatternChanges || hasWordChanges; - }, [selectedPatterns, blockedWords, originalPatterns, originalBlockedWords]); + const hasCategoryChanges = JSON.stringify(selectedContentCategories) !== JSON.stringify(originalContentCategories); + return hasPatternChanges || hasWordChanges || hasCategoryChanges; + }, [selectedPatterns, blockedWords, selectedContentCategories, originalPatterns, originalBlockedWords, originalContentCategories]); useEffect(() => { if (isEditing && onUnsavedChanges) { @@ -122,12 +166,17 @@ const ContentFilterManager: React.FC = ({ <> Content Filter Configuration {hasUnsavedChanges && ( -
-

- ⚠️ You have unsaved changes to patterns or keywords. Remember to click "Save Changes" at the - bottom. -

-
+ + You have unsaved changes to patterns or keywords. Remember to click "Save Changes" at the + bottom. + + } + /> )}
{guardrailSettings && guardrailSettings.content_filter_settings && ( @@ -150,6 +199,19 @@ const ContentFilterManager: React.FC = ({ console.log("File uploaded:", content); }} accessToken={accessToken} + contentCategories={guardrailSettings.content_filter_settings.content_categories || []} + selectedContentCategories={selectedContentCategories} + onContentCategoryAdd={(category) => + setSelectedContentCategories([...selectedContentCategories, category]) + } + onContentCategoryRemove={(id) => + setSelectedContentCategories(selectedContentCategories.filter((c) => c.id !== id)) + } + onContentCategoryUpdate={(id, field, value) => + setSelectedContentCategories( + selectedContentCategories.map((c) => (c.id === id ? { ...c, [field]: value } : c)) + ) + } /> )}
@@ -160,8 +222,16 @@ const ContentFilterManager: React.FC = ({ export default ContentFilterManager; // Helper function to format data for API -export const formatContentFilterDataForAPI = (patterns: Pattern[], blockedWords: BlockedWord[]) => { - return { +export const formatContentFilterDataForAPI = ( + patterns: Pattern[], + blockedWords: BlockedWord[], + categories?: SelectedContentCategory[] +) => { + const result: { + patterns: any[]; + blocked_words: any[]; + categories?: any[]; + } = { patterns: patterns.map((p) => ({ pattern_type: p.type === "prebuilt" ? "prebuilt" : "regex", pattern_name: p.type === "prebuilt" ? p.name : undefined, @@ -175,4 +245,13 @@ export const formatContentFilterDataForAPI = (patterns: Pattern[], blockedWords: description: w.description, })), }; + if (categories !== undefined) { + result.categories = categories.map((c) => ({ + category: c.category, + enabled: true, + action: c.action, + severity_threshold: c.severity_threshold || "medium", + })); + } + return result; }; diff --git a/ui/litellm-dashboard/src/components/guardrails/content_filter/CustomPatternModal.tsx b/ui/litellm-dashboard/src/components/guardrails/content_filter/CustomPatternModal.tsx index c69a7a0e33c..f88e5db59a4 100644 --- a/ui/litellm-dashboard/src/components/guardrails/content_filter/CustomPatternModal.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/content_filter/CustomPatternModal.tsx @@ -1,6 +1,5 @@ import React from "react"; -import { Typography, Select, Modal, Space } from "antd"; -import { Button, TextInput } from "@tremor/react"; +import { Typography, Select, Modal, Space, Button, Input } from "antd"; const { Text } = Typography; const { Option } = Select; @@ -39,20 +38,20 @@ const CustomPatternModal: React.FC = ({
Pattern name - onNameChange(e.target.value)} style={{ marginTop: 8 }} />
Regex pattern - onRegexChange(e.target.value)} style={{ marginTop: 8 }} /> @@ -77,10 +76,10 @@ const CustomPatternModal: React.FC = ({
- -
diff --git a/ui/litellm-dashboard/src/components/guardrails/content_filter/KeywordModal.tsx b/ui/litellm-dashboard/src/components/guardrails/content_filter/KeywordModal.tsx index 2a57a3d0db1..8bdbfb326d9 100644 --- a/ui/litellm-dashboard/src/components/guardrails/content_filter/KeywordModal.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/content_filter/KeywordModal.tsx @@ -1,6 +1,5 @@ import React from "react"; -import { Typography, Select, Modal, Space } from "antd"; -import { Button, TextInput, Textarea } from "@tremor/react"; +import { Typography, Select, Modal, Space, Button, Input } from "antd"; const { Text } = Typography; const { Option } = Select; @@ -39,10 +38,10 @@ const KeywordModal: React.FC = ({
Keyword - onKeywordChange(e.target.value)} style={{ marginTop: 8 }} />
@@ -64,10 +63,10 @@ const KeywordModal: React.FC = ({
Description (optional) -