scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

fix(mcp): preserve session expiry signals and scope dependency CI

Browse source
This commit is contained in:
Joshua Valluru 2026-09-18 22:52:10 -07:00
commit f5ab563499
214 changed files with 15485 additions and 1293 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
.circleci/config.yml
View file

@ -1785,6 +1785,12 @@ jobs:
- wait_for_service:
url: http://localhost:4000
timeout: "300"
- run:
name: Seed the routing strategy through /config/update
command: |
curl --noproxy '*' -sSf -X POST http://localhost:4000/config/update \
-H 'Authorization: Bearer sk-1234' -H 'Content-Type: application/json' \
-d '{"router_settings": {"routing_strategy": "usage-based-routing-v2"}}'
- run:
name: Run tests
command: |

11
.circleci/scripts/classify_changes.sh
View file

@ -1,16 +1,22 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only>}"
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only|mcp-dependencies>}"
has_client=false
has_backend=false
has_ci=false
has_provider_harness=false
has_cost_map=false
has_mcp_dependencies=false
outside_cost_map_set=false
while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
*.md | *.mdx) : ;;
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/test_litellm/test_circleci_path_filter.py | tests/test_litellm/test_detect_changes.py)
has_mcp_dependencies=true ;;
esac
case "$file" in
tests/e2e/*/*.py) : ;;
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
@ -31,6 +37,9 @@ while IFS= read -r file || [ -n "$file" ]; do
done
case "$category" in
mcp-dependencies)
[ "$has_mcp_dependencies" = true ] && echo run || echo skip
;;
cost-map-only)
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
;;

3
.circleci/scripts/run_integration.sh
View file

@ -125,6 +125,9 @@ start_proxy() {
start_proxy 4000 proxy.log
proxy_pid="$launched_pid"
.venv/bin/python .circleci/scripts/wait_integration_services.py
curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
-d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
if [ "$suite" = management ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
start_proxy 4001 peer.log

2
.github/actions/detect-changes/action.yml vendored
View file

@ -14,7 +14,7 @@ description: >-
inputs:
category:
description: "Which classification to apply: backend, client or ui"
description: "Which classification to apply: backend, client, ui, provider-harness, cost-map-only or mcp-dependencies"
required: false
default: backend
github-token:

71
.github/workflows/issue_fixed_comment.yml vendored Normal file
View file

@ -0,0 +1,71 @@
name: Issue fixed comment
on:
issues:
types: [closed]
workflow_dispatch:
inputs:
issue_number:
description: "Closed issue number to comment on manually."
required: true
pull_request:
paths:
- .github/workflows/issue_fixed_comment.yml
- scripts/comment-fixed-issue.ts
- scripts/comment-fixed-issue.test.ts
- scripts/auto-close-duplicates.ts
permissions: {}
concurrency:
group: issue-fixed-comment-${{ github.event.issue.number || github.event.inputs.issue_number || github.run_id }}
cancel-in-progress: false
jobs:
comment-fixed-issue-tests:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
bun-version: "1.4.0"
- name: Test the closer lookup, the release placement and the comment
run: bun test scripts/comment-fixed-issue.test.ts
comment-fixed-issue:
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
issues: write
steps:
- name: Checkout scripts
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: scripts
persist-credentials: false
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
# Exact version, never latest: the next step holds an issues: write token
bun-version: "1.4.0"
- name: Name the release that carries the fix
run: bun run scripts/comment-fixed-issue.ts | tee -a "${GITHUB_STEP_SUMMARY}"
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
ISSUE_NUMBER: ${{ github.event.issue.number || github.event.inputs.issue_number }}
DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}
DRY_RUN: ${{ vars.ISSUE_FIXED_COMMENT_ENABLED != 'true' }}

2
.github/workflows/test-mcp-dependency-resolution.yml vendored
View file

@ -33,6 +33,8 @@ jobs:
- name: Detect relevant changes
id: changes
uses: ./.github/actions/detect-changes
with:
category: mcp-dependencies
- name: Set up Python
if: steps.changes.outputs.decision != 'skip'

130
AGENTS.md
View file

@ -1,3 +1,131 @@
Read @CLAUDE.md for coding guidelines
Do not write comments unless they are any of:
- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)
- used as an input for tools to read and act on. For example:
- entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame
- a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # <reason>` when introducing a truly unavoidable violation
- a TODO or FIXME
- Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work
Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
- correct
- secure
- performant
- readable
- easy to maintain/change
- modern
In descending order of importance
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
Python max line length is 120, not 88
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone budget gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0`
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
Commit and push your work when you're done without asking
When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
When working on a PR, keep the PR description in sync with new commits being made
All GitHub comments must be human-readable and 15-25 words max
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>`
- Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: <reason>`
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
- Use tagged unions + match
- No monster files or god objects
- No file sprawl: deliberate file and folder structure
- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions
- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration
Follow conventional commits for commit names and PR titles
## Think Before Coding
**Don't assume. Don't hide confusion. Surface tradeoffs**
Before implementing:
- State your assumptions explicitly. If uncertain, ask
- If multiple interpretations exist, present them. Don't pick silently
- If a simpler approach exists, say so. Push back when warranted
- If something is unclear, stop. Name what's confusing. Ask
## Simplicity First
**Minimum code that solves the problem. Nothing speculative**
- No features beyond what was asked
- No abstractions for single-use code
- No "flexibility" or "configurability" that wasn't requested
- No error handling for impossible scenarios
- If you write 200 lines and it could be 50, rewrite it
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check

129
CLAUDE.md
View file

@ -1,129 +0,0 @@
Do not write comments unless they are any of:
- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)
- used as an input for tools to read and act on. For example:
- entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame
- a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # <reason>` when introducing a truly unavoidable violation
- a TODO or FIXME
- Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work
Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
- correct
- secure
- performant
- readable
- easy to maintain/change
- modern
In descending order of importance
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
Python max line length is 120, not 88
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone budget gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0`
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
Commit and push your work when you're done without asking
When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
When working on a PR, keep the PR description in sync with new commits being made
All GitHub comments must be human-readable and 15-25 words max
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>`
- Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: <reason>`
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
- Use tagged unions + match
- No monster files or god objects
- No file sprawl: deliberate file and folder structure
- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions
- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration
Follow conventional commits for commit names and PR titles
## Think Before Coding
**Don't assume. Don't hide confusion. Surface tradeoffs**
Before implementing:
- State your assumptions explicitly. If uncertain, ask
- If multiple interpretations exist, present them. Don't pick silently
- If a simpler approach exists, say so. Push back when warranted
- If something is unclear, stop. Name what's confusing. Ask
## Simplicity First
**Minimum code that solves the problem. Nothing speculative**
- No features beyond what was asked
- No abstractions for single-use code
- No "flexibility" or "configurability" that wasn't requested
- No error handling for impossible scenarios
- If you write 200 lines and it could be 50, rewrite it
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify

17
CONTRIBUTING.md
View file

@ -148,7 +148,7 @@ make lint
Individual linting commands:
```bash
make format-check # Check Black formatting
make format-check # Check ruff format formatting
make lint-ruff # Run Ruff linting
make lint-basedpyright # Run basedpyright type checking
make check-circular-imports # Check for circular imports
@ -160,14 +160,14 @@ Apply formatting (auto-fixes issues):
make format
```
> **Black formatting is enforced in CI.** All PRs must pass the Black formatting check.
> **Formatting is enforced in CI.** All PRs must pass the `ruff format --check` step.
>
> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): `AGENTS.md` and `CLAUDE.md` instruct agents to run `poetry run black .` before committing.
> - **VS Code users**: Install the [Black Formatter extension](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter) and enable format-on-save:
> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): follow `AGENTS.md` and run `make format` before committing.
> - **VS Code users**: Install the [Ruff extension](https://marketplace.visualstudio.com/items?itemName=charliermarsh.ruff) and enable format-on-save:
> ```json
> {
> "[python]": {
> "editor.defaultFormatter": "ms-python.black-formatter",
> "editor.defaultFormatter": "charliermarsh.ruff",
> "editor.formatOnSave": true
> }
> }
@ -197,8 +197,8 @@ make help # Show all available commands
make install-dev # Install development dependencies
make install-proxy-dev # Install proxy development dependencies
make install-test-deps # Install the full local test environment
make format # Apply Black code formatting
make format-check # Check Black formatting (matches CI)
make format # Apply ruff format code formatting
make format-check # Check ruff format formatting (matches CI)
make lint # Run all linting checks
make test-unit # Run unit tests
make test-integration # Run integration tests
@ -210,8 +210,7 @@ make test-unit-helm # Run Helm unit tests
LiteLLM follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html).
Our automated quality checks include:
- **Black** for consistent code formatting
- **Ruff** for linting and code quality
- **Ruff** for formatting, linting, and code quality
- **basedpyright** for static type checking
- **Circular import detection**
- **Import safety validation**

2
GEMINI.md
View file

@ -1 +1 @@
Read @CLAUDE.md for coding guidelines
Read @AGENTS.md for coding guidelines

5
README.md
View file

@ -633,9 +633,8 @@ For detailed contributing guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
LiteLLM follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html).
Our automated checks include:
- **Black** for code formatting
- **Ruff** for linting and code quality
- **MyPy** for type checking
- **Ruff** for formatting, linting, and code quality
- **basedpyright** for type checking
- **Circular import detection**
- **Import safety checks**

54
ci_cd/generate_model_prices_schema.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Optional
from typing import Final, Optional
import jsonschema
@ -19,6 +19,10 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
BOOLEAN: JsonSchema = {"type": "boolean"}
STRING: JsonSchema = {"type": "string"}
TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
WEEKDAY_PATTERN: Final = (
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
)
EXTRA_BOOLEAN_KEYS = frozenset(
{
@ -31,7 +35,51 @@ EXTRA_BOOLEAN_KEYS = frozenset(
}
)
HOURS_UTC: Final[JsonSchema] = {
"description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
"oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
}
OFF_PEAK_WINDOW: Final[JsonSchema] = {
"type": "object",
"properties": {
"hours_utc": HOURS_UTC,
"weekdays": {
"type": "array",
"description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
"items": {
"oneOf": [
{"type": "integer", "minimum": 1, "maximum": 7},
{"type": "string", "pattern": WEEKDAY_PATTERN},
]
},
"minItems": 1,
},
},
"required": ["hours_utc"],
"additionalProperties": False,
}
OBJECT_KEYS: dict[str, JsonSchema] = {
"off_peak_pricing": {
"type": "object",
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
"properties": {
"hours_utc": HOURS_UTC,
"windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
"weekday_timezone": {
"type": "string",
"description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
},
"anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
"additionalProperties": False,
},
"search_context_cost_per_query": {
"type": "object",
"description": "USD cost per web search query, keyed by search context size.",
@ -327,9 +375,7 @@ def render(schema: JsonSchema) -> str:
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
validator = jsonschema.Draft202012Validator(
schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER
)
validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
return tuple(
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)

43
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -11,7 +11,7 @@ Endpoints for /project operations
#### PROJECT MANAGEMENT ####
import json
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Request
@ -22,7 +22,11 @@ from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import delete_cached_project_object
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
from litellm.proxy.management_endpoints.common_utils import (
_is_user_team_admin, # pyright: ignore[reportPrivateUsage] # shared owner of team-admin membership
_set_object_metadata_field,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
@ -82,37 +86,38 @@ async def _check_user_permission_for_project(
user_api_key_dict: UserAPIKeyAuth,
team_id: str | None,
prisma_client: PrismaClient,
general_settings: Mapping[str, object],
require_admin: bool = False,
team_object: LiteLLM_TeamTable | None = None,
) -> bool:
"""
Check if user has permission to manage a project.
Returns True if user is proxy admin or team admin (when team_id provided).
Returns True if user is proxy admin, or a team admin of ``team_id`` when the
``team_admin_editable_team_fields`` setting grants team admins the ``projects`` permission.
If require_admin=True, only proxy admins are allowed.
If team_object is provided, it will be used instead of fetching from DB
(avoids duplicate DB queries when team was already fetched for validation).
"""
is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
if require_admin:
if require_admin or is_proxy_admin:
return is_proxy_admin
if is_proxy_admin:
return True
if not team_id or not user_api_key_dict.user_id:
if not team_id or not user_api_key_dict.user_id or not team_admin_may_manage_projects(general_settings):
return False
team = team_object
if team is None:
team = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
team_row: Final = (
team_object
if team_object is not None
else await _team_table(prisma_client).find_unique(where={"team_id": team_id})
)
if team_row is None:
return False
if team and team.admins:
return user_api_key_dict.user_id in team.admins
return False
team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump())
return _is_user_team_admin(user_api_key_dict, team) or user_api_key_dict.user_id in (team.admins or [])
async def _validate_team_exists(
@ -531,6 +536,7 @@ async def new_project(
user_api_key_dict=user_api_key_dict,
team_id=data.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()),
)
@ -735,6 +741,7 @@ async def update_project(
user_api_key_dict=user_api_key_dict,
team_id=existing_project.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
)
if not has_permission:
@ -751,6 +758,7 @@ async def update_project(
user_api_key_dict=user_api_key_dict,
team_id=data.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
team_object=(
LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()) if target_team_obj else None
),
@ -877,7 +885,7 @@ async def delete_project(
}'
```
"""
from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import general_settings, premium_user, prisma_client, user_api_key_cache
try:
if not premium_user:
@ -899,6 +907,7 @@ async def delete_project(
user_api_key_dict=user_api_key_dict,
team_id=None,
prisma_client=prisma_client,
general_settings=general_settings,
require_admin=True,
)

19
litellm-rust/Cargo.lock generated
View file

@ -2050,8 +2050,10 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-host",
"litellm-http",
"litellm-llms",
"litellm-types",
"mime_guess",
@ -2129,6 +2131,20 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-http"
version = "0.1.0"
dependencies = [
"http 1.4.2",
"hyper-util",
"reqwest 0.12.28",
"rstest",
"rustls 0.23.42",
"thiserror 2.0.19",
"tokio",
"webpki-roots",
]
[[package]]
name = "litellm-llms"
version = "0.1.0"
@ -2146,6 +2162,7 @@ dependencies = [
"litellm-core-utils",
"litellm-framing",
"litellm-host",
"litellm-http",
"litellm-types",
"reqwest 0.12.28",
"rstest",
@ -2167,9 +2184,11 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-auth-gcp",
"litellm-callbacks-legacy",
"litellm-core",
"litellm-host-python",
"litellm-http",
"litellm-llms",
"litellm-token-counter",
"litellm-types",

4
litellm-rust/Cargo.toml
View file

@ -17,6 +17,7 @@ litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
litellm-auth-gcp = { path = "crates/auth-gcp" }
litellm-http = { path = "crates/http" }
litellm-llms = { path = "crates/llms" }
litellm-types = { path = "crates/types" }
litellm-core-utils = { path = "crates/core-utils" }
@ -26,6 +27,8 @@ litellm-token-counter = { path = "crates/token-counter" }
litellm-host-python = { path = "crates/host-python" }
bytes = "1"
http = "1"
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
proptest = "1.7.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
@ -49,6 +52,7 @@ base64 = "0.22"
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
webpki-roots = "1"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
fancy-regex = "0.19.2"

2
litellm-rust/crates/core/Cargo.toml
View file

@ -35,6 +35,8 @@ url.workspace = true
veil.workspace = true
[dev-dependencies]
litellm-auth-gcp.workspace = true
litellm-http.workspace = true
litellm-llms = { workspace = true, features = ["test-support"] }
rstest.workspace = true
rstest_reuse.workspace = true

4
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -14,7 +14,3 @@ pub async fn perform(
) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
perform(&OcrClient::shared()?, request).await
}

36
litellm-rust/crates/core/tests/ocr.rs
View file

@ -1,16 +1,21 @@
use std::sync::{Arc, Mutex};
use litellm_auth_gcp::VertexAuth;
use litellm_host::{
event::{CallEvent, MachineEvent, WireRequest},
host::{Host, HostOp, HostResult},
machine::{HostFailure, Machine, MachineStep},
};
use litellm_http::{HttpClientPool, HttpSettings, Resolution};
use litellm_llms::{
base_llm::ocr::{
error::Error as OcrError,
transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig},
},
custom_httpx::llm_http_handler::OcrClient,
custom_httpx::{
llm_http_handler::OcrClient,
media::{PublicDnsResolver, UrlPolicy},
},
};
use rstest::rstest;
use serde_json::{Value, json};
@ -171,25 +176,24 @@ async fn facade_retains_native_response_when_requested() {
}
#[tokio::test]
async fn facade_uses_the_injected_http_client() {
async fn ocr_client_uses_the_injected_http_pool_configuration() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let mut default_headers = reqwest::header::HeaderMap::new();
default_headers.insert(
"x-transport-owner",
reqwest::header::HeaderValue::from_static("host"),
);
let provider_http = reqwest::Client::builder()
.default_headers(default_headers)
.build()
.unwrap();
crate::ocr::client::perform(
&OcrClient::new(provider_http).unwrap(),
wire_request("mistral/model", &base, json!({})),
let settings = HttpSettings {
user_agent: Some("host-owned/1".into()),
..HttpSettings::default()
};
let client = OcrClient::new(
&HttpClientPool::new(Arc::new(PublicDnsResolver)),
&Resolution::from(&settings).config,
UrlPolicy::default(),
VertexAuth::default(),
)
.await
.unwrap();
crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({})))
.await
.unwrap();
server.await.unwrap();
assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host"));
assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1"));
}
fn event_name(event: &CallEvent) -> &'static str {

1
litellm-rust/crates/http/AGENTS.md Normal file
View file

@ -0,0 +1 @@
- https://github.com/BerriAI/litellm-docs/blob/main/docs/guides/security_settings.md

18
litellm-rust/crates/http/Cargo.toml Normal file
View file

@ -0,0 +1,18 @@
[package]
name = "litellm-http"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
http.workspace = true
hyper-util.workspace = true
reqwest.workspace = true
rustls.workspace = true
thiserror.workspace = true
webpki-roots.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

296
litellm-rust/crates/http/src/config.rs Normal file
View file

@ -0,0 +1,296 @@
use std::{
net::{IpAddr, Ipv4Addr},
path::PathBuf,
time::Duration,
};
use crate::{
error::Error,
settings::{HttpSettings, SslVerify, TcpKeepalive},
tls::{CipherSelection, KeyExchangeGroup, Tls12CipherSuite, Unsupported},
};
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum Verify {
Disabled,
CaBundle(PathBuf),
BuiltInRoots,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct HttpClientConfig {
pub verify: Verify,
pub client_certificate: Option<PathBuf>,
pub key_exchange_group: Option<KeyExchangeGroup>,
pub tls12_cipher_suites: Option<Vec<Tls12CipherSuite>>,
pub force_ipv4: bool,
pub http2: bool,
pub user_agent: Option<String>,
pub trust_proxy_env: bool,
pub connect_timeout: Duration,
pub tcp_keepalive: Option<TcpKeepalive>,
pub pool_idle_timeout: Duration,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Resolution {
pub config: HttpClientConfig,
pub unsupported: Vec<Unsupported>,
}
impl From<&HttpSettings> for Verify {
fn from(settings: &HttpSettings) -> Self {
match &settings.ssl_verify {
Some(SslVerify::Disabled) => Self::Disabled,
Some(SslVerify::CaBundle(path)) => Self::CaBundle(path.clone()),
Some(SslVerify::Enabled) | None => settings
.ssl_cert_file
.clone()
.map_or(Self::BuiltInRoots, Self::CaBundle),
}
}
}
impl From<&HttpSettings> for Resolution {
fn from(settings: &HttpSettings) -> Self {
let curve = settings
.ssl_ecdh_curve
.as_deref()
.map(str::parse::<KeyExchangeGroup>)
.transpose();
let ciphers = settings
.ssl_security_level
.as_deref()
.map(CipherSelection::from)
.unwrap_or_default();
Self {
config: HttpClientConfig {
verify: Verify::from(settings),
client_certificate: settings.ssl_certificate.clone(),
key_exchange_group: curve.clone().ok().flatten(),
tls12_cipher_suites: ciphers.tls12_cipher_suites,
force_ipv4: settings.force_ipv4,
http2: settings.http2,
user_agent: settings.user_agent.clone(),
trust_proxy_env: settings.trust_proxy_env,
connect_timeout: settings.connect_timeout,
tcp_keepalive: settings.tcp_keepalive,
pool_idle_timeout: settings.pool_idle_timeout,
},
unsupported: curve.err().into_iter().chain(ciphers.unsupported).collect(),
}
}
}
impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder {
type Error = Error;
fn try_from(config: &HttpClientConfig) -> Result<Self, Self::Error> {
let base = reqwest::Client::builder()
.use_preconfigured_tls(rustls::ClientConfig::try_from(config)?)
.connect_timeout(config.connect_timeout)
.pool_idle_timeout(config.pool_idle_timeout);
let with_keepalive = match config.tcp_keepalive {
None => base,
Some(keepalive) => base
.tcp_keepalive(keepalive.idle)
.tcp_keepalive_interval(keepalive.interval)
.tcp_keepalive_retries(keepalive.retries),
};
let with_address = if config.force_ipv4 {
with_keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
} else {
with_keepalive
};
let with_protocol = if config.http2 {
with_address
} else {
with_address.http1_only()
};
let with_agent = match &config.user_agent {
Some(agent) => with_protocol.user_agent(agent),
None => with_protocol,
};
Ok(if config.trust_proxy_env {
with_agent
} else {
with_agent.no_proxy()
})
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn settings(ssl_verify: Option<SslVerify>, ssl_cert_file: Option<&str>) -> HttpSettings {
HttpSettings {
ssl_verify,
ssl_cert_file: ssl_cert_file.map(PathBuf::from),
..HttpSettings::default()
}
}
#[rstest]
#[case::default(settings(None, None), Verify::BuiltInRoots)]
#[case::setting_disables(
settings(Some(SslVerify::Disabled), Some("/env/roots.pem")),
Verify::Disabled
)]
#[case::setting_bundle(
settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")),
Verify::CaBundle("/configured.pem".into())
)]
#[case::enabled_uses_cert_file(
settings(Some(SslVerify::Enabled), Some("/env/roots.pem")),
Verify::CaBundle("/env/roots.pem".into())
)]
#[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), Verify::CaBundle("/env/roots.pem".into()))]
fn verify_follows_setting_then_cert_file(
#[case] settings: HttpSettings,
#[case] expected: Verify,
) {
let config = Resolution::from(&settings).config;
assert_eq!(config.verify, expected);
}
#[rstest]
#[case::x25519("X25519", Some(KeyExchangeGroup::X25519))]
#[case::openssl_p256("prime256v1", Some(KeyExchangeGroup::Secp256r1))]
#[case::p384("secp384r1", Some(KeyExchangeGroup::Secp384r1))]
fn ecdh_curve_selects_the_single_key_exchange_group(
#[case] curve: &str,
#[case] expected: Option<KeyExchangeGroup>,
) {
let settings = HttpSettings {
ssl_ecdh_curve: Some(curve.into()),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(resolution.config.key_exchange_group, expected);
assert_eq!(resolution.unsupported, []);
}
#[test]
fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported() {
let settings = HttpSettings {
ssl_ecdh_curve: Some("secp521r1".into()),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(resolution.config.key_exchange_group, None);
assert_eq!(
resolution.unsupported,
[Unsupported::EcdhCurve("secp521r1".into())]
);
}
#[test]
fn legacy_security_level_keeps_every_suite_and_is_reported_unsupported() {
let settings = HttpSettings {
ssl_security_level: Some("DEFAULT@SECLEVEL=1".into()),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(resolution.config.tls12_cipher_suites, None);
assert_eq!(
resolution.unsupported,
[Unsupported::SecurityLevel("@SECLEVEL=1".into())]
);
}
#[test]
fn named_suites_restrict_tls12_and_unsupported_entries_are_reported() {
let settings = HttpSettings {
ssl_security_level: Some(
"ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES128-GCM-SHA256:!aNULL:AES256-SHA@SECLEVEL=2"
.into(),
),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(
resolution.config.tls12_cipher_suites,
Some(vec![
Tls12CipherSuite::EcdheEcdsaAes128Gcm,
Tls12CipherSuite::EcdheRsaAes256Gcm
])
);
assert_eq!(
resolution.unsupported,
[
Unsupported::CipherToken("!aNULL".into()),
Unsupported::CipherToken("AES256-SHA".into())
]
);
}
#[test]
fn connection_settings_carry_over_unchanged() {
let keepalive = TcpKeepalive {
idle: Duration::from_secs(60),
interval: Duration::from_secs(30),
retries: 5,
};
let settings = HttpSettings {
ssl_certificate: Some("/client.pem".into()),
force_ipv4: true,
http2: true,
user_agent: Some("litellm/1.0".into()),
trust_proxy_env: true,
connect_timeout: Duration::from_secs(7),
tcp_keepalive: Some(keepalive),
pool_idle_timeout: Duration::from_secs(45),
..HttpSettings::default()
};
let config = Resolution::from(&settings).config;
assert_eq!(
config,
HttpClientConfig {
verify: Verify::BuiltInRoots,
client_certificate: Some("/client.pem".into()),
key_exchange_group: None,
tls12_cipher_suites: None,
force_ipv4: true,
http2: true,
user_agent: Some("litellm/1.0".into()),
trust_proxy_env: true,
connect_timeout: Duration::from_secs(7),
tcp_keepalive: Some(keepalive),
pool_idle_timeout: Duration::from_secs(45),
}
);
}
#[test]
fn missing_ca_bundle_is_a_read_error() {
let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem");
let config = HttpClientConfig {
verify: Verify::CaBundle(path.clone()),
..Resolution::from(&HttpSettings::default()).config
};
assert!(matches!(
reqwest::ClientBuilder::try_from(&config),
Err(Error::Read { path: reported, .. }) if reported == path
));
}
#[test]
fn non_pem_ca_bundle_is_an_invalid_pem_error() {
let path =
std::env::temp_dir().join(format!("litellm-http-not-pem-{}.pem", std::process::id()));
std::fs::write(&path, b"not a certificate").unwrap();
let config = HttpClientConfig {
verify: Verify::CaBundle(path.clone()),
..Resolution::from(&HttpSettings::default()).config
};
let result = reqwest::ClientBuilder::try_from(&config).map(drop);
std::fs::remove_file(&path).unwrap();
assert!(matches!(
result,
Err(Error::InvalidPem { path: reported, .. }) if reported == path
));
}
}

17
litellm-rust/crates/http/src/error.rs Normal file
View file

@ -0,0 +1,17 @@
use std::path::PathBuf;
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("could not read {}: {message}", path.display())]
Read { path: PathBuf, message: String },
#[error("{} is not a PEM file: {message}", path.display())]
InvalidPem { path: PathBuf, message: String },
#[error("could not build the HTTP client: {0}")]
Client(String),
}
impl From<reqwest::Error> for Error {
fn from(error: reqwest::Error) -> Self {
Self::Client(error.without_url().to_string())
}
}

13
litellm-rust/crates/http/src/lib.rs Normal file
View file

@ -0,0 +1,13 @@
mod config;
mod error;
mod pool;
mod proxy;
mod settings;
mod tls;
pub use config::{HttpClientConfig, Resolution, Verify};
pub use error::Error;
pub use pool::{ClientVariant, HttpClientPool};
pub use proxy::EnvironmentProxies;
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};
pub use tls::{KeyExchangeGroup, Tls12CipherSuite, Unsupported};

325
litellm-rust/crates/http/src/pool.rs Normal file
View file

@ -0,0 +1,325 @@
use std::{
collections::HashMap,
sync::{Arc, Mutex, MutexGuard, PoisonError},
time::{Duration, Instant},
};
use reqwest::dns::Resolve;
use crate::{config::HttpClientConfig, error::Error};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ClientVariant {
Provider,
NoRedirect,
Media,
UnpinnedMedia,
}
const CLIENT_TTL: Duration = Duration::from_secs(3600);
struct PooledClient {
client: reqwest::Client,
built_at: Instant,
}
type Clients = HashMap<(HttpClientConfig, ClientVariant), PooledClient>;
pub struct HttpClientPool {
media_resolver: Arc<dyn Resolve>,
ttl: Duration,
clients: Mutex<Clients>,
}
impl HttpClientPool {
pub fn new(media_resolver: Arc<dyn Resolve>) -> Self {
Self::with_ttl(media_resolver, CLIENT_TTL)
}
pub fn with_ttl(media_resolver: Arc<dyn Resolve>, ttl: Duration) -> Self {
Self {
media_resolver,
ttl,
clients: Mutex::default(),
}
}
pub fn client(
&self,
config: &HttpClientConfig,
variant: ClientVariant,
) -> Result<reqwest::Client, Error> {
let effective = match variant {
ClientVariant::Media => HttpClientConfig {
client_certificate: None,
trust_proxy_env: false,
..config.clone()
},
ClientVariant::UnpinnedMedia => HttpClientConfig {
client_certificate: None,
..config.clone()
},
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
};
let key = (effective, variant);
if let Some(pooled) = self.lock().get(&key)
&& pooled.built_at.elapsed() < self.ttl
{
return Ok(pooled.client.clone());
}
let client = self
.apply(variant, reqwest::ClientBuilder::try_from(&key.0)?)
.build()?;
self.lock().insert(
key,
PooledClient {
client: client.clone(),
built_at: Instant::now(),
},
);
Ok(client)
}
fn lock(&self) -> MutexGuard<'_, Clients> {
self.clients.lock().unwrap_or_else(PoisonError::into_inner)
}
fn apply(
&self,
variant: ClientVariant,
builder: reqwest::ClientBuilder,
) -> reqwest::ClientBuilder {
match variant {
ClientVariant::Provider => builder,
ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => {
builder.redirect(reqwest::redirect::Policy::none())
}
ClientVariant::Media => builder
.redirect(reqwest::redirect::Policy::none())
.dns_resolver2(Arc::clone(&self.media_resolver)),
}
}
}
#[cfg(test)]
mod tests {
use std::{
net::SocketAddr,
sync::atomic::{AtomicUsize, Ordering},
time::Duration,
};
use reqwest::dns::{Addrs, Name, Resolving};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
use super::*;
use crate::{HttpSettings, Resolution, Verify};
struct FixedResolver(SocketAddr);
impl Resolve for FixedResolver {
fn resolve(&self, _: Name) -> Resolving {
let addrs: Addrs = Box::new(std::iter::once(self.0));
Box::pin(std::future::ready(Ok(addrs)))
}
}
fn pool() -> HttpClientPool {
HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())))
}
fn config(user_agent: &str) -> HttpClientConfig {
HttpClientConfig {
user_agent: Some(user_agent.into()),
..Resolution::from(&HttpSettings::default()).config
}
}
async fn serve(
status_line: &'static str,
) -> (SocketAddr, Arc<AtomicUsize>, Arc<Mutex<Vec<String>>>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let connections = Arc::new(AtomicUsize::new(0));
let requests = Arc::new(Mutex::new(Vec::new()));
let (accepted, seen) = (Arc::clone(&connections), Arc::clone(&requests));
tokio::spawn(async move {
loop {
let (mut socket, _) = listener.accept().await.unwrap();
accepted.fetch_add(1, Ordering::SeqCst);
let seen = Arc::clone(&seen);
tokio::spawn(async move {
let mut buffer = vec![0u8; 4096];
while let Ok(read) = socket.read(&mut buffer).await {
if read == 0 {
return;
}
seen.lock()
.unwrap()
.push(String::from_utf8_lossy(&buffer[..read]).into_owned());
let response = format!(
"{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\n\r\n"
);
if socket.write_all(response.as_bytes()).await.is_err() {
return;
}
}
});
}
});
(address, connections, requests)
}
async fn get(
pool: &HttpClientPool,
config: &HttpClientConfig,
variant: ClientVariant,
url: &str,
) -> reqwest::Response {
pool.client(config, variant)
.unwrap()
.get(url)
.timeout(Duration::from_secs(5))
.send()
.await
.unwrap()
}
#[tokio::test]
async fn clients_are_shared_per_config_and_variant() {
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
let url = format!("http://{address}");
let pool = pool();
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
assert_eq!(connections.load(Ordering::SeqCst), 1);
get(&pool, &config("a"), ClientVariant::NoRedirect, &url).await;
assert_eq!(connections.load(Ordering::SeqCst), 2);
get(&pool, &config("b"), ClientVariant::Provider, &url).await;
assert_eq!(connections.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn expired_clients_are_rebuilt() {
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
let url = format!("http://{address}");
let pool = HttpClientPool::with_ttl(
Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())),
Duration::ZERO,
);
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
assert_eq!(connections.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn media_clients_are_shared_across_proxy_settings_they_never_use() {
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
let pool = HttpClientPool::new(Arc::new(FixedResolver(address)));
let url = format!("http://media.invalid:{}/doc", address.port());
for trust_proxy_env in [true, false] {
let config = HttpClientConfig {
trust_proxy_env,
..config("a")
};
get(&pool, &config, ClientVariant::Media, &url).await;
}
assert_eq!(connections.load(Ordering::SeqCst), 1);
}
#[test]
fn media_variant_never_loads_the_client_certificate() {
let pool = pool();
let with_identity = HttpClientConfig {
client_certificate: Some(std::env::temp_dir().join("litellm-http-absent-client.pem")),
..config("a")
};
assert!(
pool.client(&with_identity, ClientVariant::Provider)
.is_err()
);
assert!(pool.client(&with_identity, ClientVariant::Media).is_ok());
assert!(
pool.client(&with_identity, ClientVariant::UnpinnedMedia)
.is_ok()
);
}
#[test]
fn build_failures_are_not_cached() {
let pool = pool();
let missing = HttpClientConfig {
verify: Verify::CaBundle(std::env::temp_dir().join("litellm-http-absent.pem")),
..config("a")
};
assert!(pool.client(&missing, ClientVariant::Provider).is_err());
assert!(pool.client(&missing, ClientVariant::Provider).is_err());
assert!(pool.client(&config("a"), ClientVariant::Provider).is_ok());
}
#[tokio::test]
async fn provider_client_sends_the_configured_user_agent_over_http1() {
let (address, _, requests) = serve("HTTP/1.1 204 No Content").await;
let response = get(
&pool(),
&config("litellm-test/9"),
ClientVariant::Provider,
&format!("http://{address}"),
)
.await;
assert_eq!(response.status(), 204);
assert_eq!(response.version(), reqwest::Version::HTTP_11);
let request = requests.lock().unwrap()[0].clone();
assert!(request.contains("user-agent: litellm-test/9"), "{request}");
}
#[tokio::test]
async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() {
let (address, _, _) = serve("HTTP/1.1 302 Found").await;
let response = get(
&pool(),
&config("a"),
ClientVariant::NoRedirect,
&format!("http://{address}"),
)
.await;
assert_eq!(response.status(), 302);
assert_eq!(response.headers()["location"], "/elsewhere");
}
#[tokio::test]
async fn unpinned_media_variant_uses_the_system_resolver_and_returns_redirects() {
let (address, _, _) = serve("HTTP/1.1 302 Found").await;
let pool = HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())));
let response = get(
&pool,
&config("a"),
ClientVariant::UnpinnedMedia,
&format!("http://localhost:{}/doc", address.port()),
)
.await;
assert_eq!(response.status(), 302);
}
#[tokio::test]
async fn media_variant_resolves_through_the_injected_resolver() {
let (address, _, requests) = serve("HTTP/1.1 204 No Content").await;
let pool = HttpClientPool::new(Arc::new(FixedResolver(address)));
let url = format!("http://media.invalid:{}/doc", address.port());
let response = get(&pool, &config("a"), ClientVariant::Media, &url).await;
assert_eq!(response.status(), 204);
assert!(requests.lock().unwrap()[0].contains("host: media.invalid"));
assert!(
pool.client(&config("a"), ClientVariant::Provider)
.unwrap()
.get(&url)
.timeout(Duration::from_secs(5))
.send()
.await
.is_err()
);
}
}

15
litellm-rust/crates/http/src/proxy.rs Normal file
View file

@ -0,0 +1,15 @@
use hyper_util::client::proxy::matcher::Matcher;
pub struct EnvironmentProxies(Matcher);
impl EnvironmentProxies {
pub fn from_environment() -> Self {
Self(Matcher::from_system())
}
pub fn apply_to(&self, url: &reqwest::Url) -> bool {
url.as_str()
.parse::<http::Uri>()
.is_ok_and(|uri| self.0.intercept(&uri).is_some())
}
}

416
litellm-rust/crates/http/src/settings.rs Normal file
View file

@ -0,0 +1,416 @@
use std::{
path::{Path, PathBuf},
time::Duration,
};
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum SslVerify {
Enabled,
Disabled,
CaBundle(PathBuf),
}
impl SslVerify {
pub fn parse(value: &str) -> Self {
match value.trim().to_ascii_lowercase().as_str() {
"true" => Self::Enabled,
"false" => Self::Disabled,
_ => Self::CaBundle(PathBuf::from(value)),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct TcpKeepalive {
pub idle: Duration,
pub interval: Duration,
pub retries: u32,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct HttpSettingsLayer {
pub ssl_verify: Option<SslVerify>,
pub ssl_cert_file: Option<PathBuf>,
pub ssl_certificate: Option<PathBuf>,
pub ssl_security_level: Option<String>,
pub ssl_ecdh_curve: Option<String>,
pub force_ipv4: Option<bool>,
pub http2: Option<bool>,
pub aiohttp_trust_env: Option<bool>,
pub disable_aiohttp_trust_env: Option<bool>,
pub disable_aiohttp_transport: Option<bool>,
pub user_agent: Option<String>,
pub tcp_keepalive: Option<TcpKeepalive>,
pub pool_idle_timeout: Option<Duration>,
}
impl HttpSettingsLayer {
pub fn from_environment(env: &(dyn Fn(&str) -> Option<String> + Sync)) -> Self {
let enabled = |name: &str| {
env(name)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
.then_some(true)
};
let number = |name: &str| env(name).and_then(|value| value.trim().parse::<u32>().ok());
let seconds = |name: &str, default: u32| {
Duration::from_secs(u64::from(number(name).unwrap_or(default)))
};
Self {
ssl_verify: env("SSL_VERIFY").map(|value| SslVerify::parse(&value)),
ssl_cert_file: env("SSL_CERT_FILE").map(PathBuf::from),
ssl_certificate: env("SSL_CERTIFICATE").map(PathBuf::from),
ssl_security_level: env("SSL_SECURITY_LEVEL"),
ssl_ecdh_curve: env("SSL_ECDH_CURVE"),
force_ipv4: None,
http2: enabled("LITELLM_HTTP2"),
aiohttp_trust_env: enabled("AIOHTTP_TRUST_ENV"),
disable_aiohttp_trust_env: enabled("DISABLE_AIOHTTP_TRUST_ENV"),
disable_aiohttp_transport: enabled("DISABLE_AIOHTTP_TRANSPORT"),
user_agent: env("LITELLM_USER_AGENT"),
tcp_keepalive: enabled("AIOHTTP_SO_KEEPALIVE").map(|_| TcpKeepalive {
idle: seconds("AIOHTTP_TCP_KEEPIDLE", 60),
interval: seconds("AIOHTTP_TCP_KEEPINTVL", 30),
retries: number("AIOHTTP_TCP_KEEPCNT").unwrap_or(5),
}),
pool_idle_timeout: number("AIOHTTP_KEEPALIVE_TIMEOUT")
.map(|timeout| Duration::from_secs(u64::from(timeout))),
}
}
fn or(self, lower: Self) -> Self {
Self {
ssl_verify: self.ssl_verify.or(lower.ssl_verify),
ssl_cert_file: self.ssl_cert_file.or(lower.ssl_cert_file),
ssl_certificate: self.ssl_certificate.or(lower.ssl_certificate),
ssl_security_level: self.ssl_security_level.or(lower.ssl_security_level),
ssl_ecdh_curve: self.ssl_ecdh_curve.or(lower.ssl_ecdh_curve),
force_ipv4: self.force_ipv4.or(lower.force_ipv4),
http2: self.http2.or(lower.http2),
aiohttp_trust_env: self.aiohttp_trust_env.or(lower.aiohttp_trust_env),
disable_aiohttp_trust_env: self
.disable_aiohttp_trust_env
.or(lower.disable_aiohttp_trust_env),
disable_aiohttp_transport: self
.disable_aiohttp_transport
.or(lower.disable_aiohttp_transport),
user_agent: self.user_agent.or(lower.user_agent),
tcp_keepalive: self.tcp_keepalive.or(lower.tcp_keepalive),
pool_idle_timeout: self.pool_idle_timeout.or(lower.pool_idle_timeout),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HttpSettings {
pub ssl_verify: Option<SslVerify>,
pub ssl_cert_file: Option<PathBuf>,
pub ssl_certificate: Option<PathBuf>,
pub ssl_security_level: Option<String>,
pub ssl_ecdh_curve: Option<String>,
pub force_ipv4: bool,
pub http2: bool,
pub user_agent: Option<String>,
pub trust_proxy_env: bool,
pub connect_timeout: Duration,
pub tcp_keepalive: Option<TcpKeepalive>,
pub pool_idle_timeout: Duration,
}
impl Default for HttpSettings {
fn default() -> Self {
Self {
ssl_verify: None,
ssl_cert_file: None,
ssl_certificate: None,
ssl_security_level: None,
ssl_ecdh_curve: None,
force_ipv4: false,
http2: false,
user_agent: None,
trust_proxy_env: true,
connect_timeout: Duration::from_secs(10),
tcp_keepalive: None,
pool_idle_timeout: Duration::from_secs(120),
}
}
}
impl HttpSettings {
pub fn from_layers(
highest_precedence_first: impl IntoIterator<Item = HttpSettingsLayer>,
) -> Self {
let merged = highest_precedence_first
.into_iter()
.reduce(HttpSettingsLayer::or)
.unwrap_or_default();
let defaults = Self::default();
let http2 = merged.http2.unwrap_or(defaults.http2);
Self {
ssl_verify: merged.ssl_verify,
ssl_cert_file: merged.ssl_cert_file,
ssl_certificate: merged
.ssl_certificate
.filter(|path| !path.as_os_str().is_empty()),
ssl_security_level: merged.ssl_security_level.filter(|level| !level.is_empty()),
ssl_ecdh_curve: merged.ssl_ecdh_curve.filter(|curve| !curve.is_empty()),
force_ipv4: merged.force_ipv4.unwrap_or(defaults.force_ipv4),
http2,
user_agent: merged.user_agent,
trust_proxy_env: !merged.disable_aiohttp_trust_env.unwrap_or(false)
|| merged.aiohttp_trust_env.unwrap_or(false)
|| merged.disable_aiohttp_transport.unwrap_or(false)
|| http2,
tcp_keepalive: merged.tcp_keepalive,
pool_idle_timeout: merged
.pool_idle_timeout
.unwrap_or(defaults.pool_idle_timeout),
..defaults
}
}
pub fn without_missing_files(self, exists: &dyn Fn(&Path) -> bool) -> Self {
Self {
ssl_verify: match self.ssl_verify {
Some(SslVerify::CaBundle(path)) if !exists(&path) => Some(SslVerify::Enabled),
other => other,
},
ssl_cert_file: self.ssl_cert_file.filter(|path| exists(path)),
..self
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn no_env(_: &str) -> Option<String> {
None
}
fn env_of(
values: &'static [(&'static str, &'static str)],
) -> impl Fn(&str) -> Option<String> + Sync {
move |name| {
values
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| value.to_string())
}
}
#[rstest]
#[case("true", SslVerify::Enabled)]
#[case(" True ", SslVerify::Enabled)]
#[case("FALSE", SslVerify::Disabled)]
#[case("/etc/ssl/bundle.pem", SslVerify::CaBundle("/etc/ssl/bundle.pem".into()))]
fn ssl_verify_parses_bools_and_treats_anything_else_as_a_bundle_path(
#[case] value: &str,
#[case] expected: SslVerify,
) {
assert_eq!(SslVerify::parse(value), expected);
}
#[test]
fn higher_layers_override_lower_ones() {
let configured = HttpSettingsLayer {
ssl_verify: Some(SslVerify::Enabled),
ssl_certificate: Some("/configured/client.pem".into()),
ssl_security_level: Some("configured".into()),
user_agent: Some("configured/1".into()),
..HttpSettingsLayer::default()
};
let environment = HttpSettingsLayer::from_environment(&env_of(&[
("SSL_VERIFY", "false"),
("SSL_CERT_FILE", "/env/roots.pem"),
("SSL_CERTIFICATE", "/env/client.pem"),
("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1"),
("SSL_ECDH_CURVE", "X25519"),
("LITELLM_USER_AGENT", "env/2"),
]));
let settings = HttpSettings::from_layers([environment, configured]);
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
assert_eq!(settings.ssl_cert_file, Some("/env/roots.pem".into()));
assert_eq!(settings.ssl_certificate, Some("/env/client.pem".into()));
assert_eq!(
settings.ssl_security_level.as_deref(),
Some("DEFAULT@SECLEVEL=1")
);
assert_eq!(settings.ssl_ecdh_curve.as_deref(), Some("X25519"));
assert_eq!(settings.user_agent.as_deref(), Some("env/2"));
}
#[test]
fn an_explicit_false_in_a_higher_layer_beats_a_lower_true() {
let higher = HttpSettingsLayer {
http2: Some(false),
force_ipv4: Some(false),
..HttpSettingsLayer::default()
};
let lower = HttpSettingsLayer {
http2: Some(true),
force_ipv4: Some(true),
..HttpSettingsLayer::default()
};
let settings = HttpSettings::from_layers([higher, lower]);
assert!(!settings.http2);
assert!(!settings.force_ipv4);
}
#[test]
fn an_empty_environment_is_an_empty_layer_so_lower_layers_and_defaults_apply() {
assert_eq!(
HttpSettingsLayer::from_environment(&no_env),
HttpSettingsLayer::default()
);
let configured = HttpSettingsLayer {
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
http2: Some(true),
user_agent: Some("configured/1".into()),
..HttpSettingsLayer::default()
};
assert_eq!(
HttpSettings::from_layers([HttpSettingsLayer::default(), configured]),
HttpSettings {
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
http2: true,
user_agent: Some("configured/1".into()),
..HttpSettings::default()
}
);
assert_eq!(HttpSettings::from_layers([]), HttpSettings::default());
}
#[test]
fn empty_environment_values_clear_the_setting_like_python_truthiness() {
let configured = HttpSettingsLayer {
ssl_certificate: Some("/configured/client.pem".into()),
ssl_security_level: Some("configured".into()),
ssl_ecdh_curve: Some("X25519".into()),
..HttpSettingsLayer::default()
};
let environment = HttpSettingsLayer::from_environment(&env_of(&[
("SSL_CERTIFICATE", ""),
("SSL_SECURITY_LEVEL", ""),
("SSL_ECDH_CURVE", ""),
]));
let settings = HttpSettings::from_layers([environment, configured]);
assert_eq!(settings.ssl_certificate, None);
assert_eq!(settings.ssl_security_level, None);
assert_eq!(settings.ssl_ecdh_curve, None);
}
#[test]
fn socket_keepalive_follows_the_aiohttp_variables_with_python_defaults() {
let tuned = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(&[
("AIOHTTP_SO_KEEPALIVE", "True"),
("AIOHTTP_TCP_KEEPIDLE", "45"),
("AIOHTTP_KEEPALIVE_TIMEOUT", "30"),
]))]);
assert_eq!(
tuned.tcp_keepalive,
Some(TcpKeepalive {
idle: Duration::from_secs(45),
interval: Duration::from_secs(30),
retries: 5,
})
);
assert_eq!(tuned.pool_idle_timeout, Duration::from_secs(30));
}
#[test]
fn socket_keepalive_stays_off_unless_enabled() {
let settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(
&[("AIOHTTP_TCP_KEEPIDLE", "45")],
))]);
assert_eq!(settings.tcp_keepalive, None);
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
}
fn proxy_flags(
aiohttp_trust_env: bool,
disable_aiohttp_trust_env: bool,
disable_aiohttp_transport: bool,
http2: bool,
) -> HttpSettingsLayer {
HttpSettingsLayer {
aiohttp_trust_env: Some(aiohttp_trust_env),
disable_aiohttp_trust_env: Some(disable_aiohttp_trust_env),
disable_aiohttp_transport: Some(disable_aiohttp_transport),
http2: Some(http2),
..HttpSettingsLayer::default()
}
}
#[rstest]
#[case::aiohttp_default(proxy_flags(false, false, false, false), true)]
#[case::aiohttp_opted_out(proxy_flags(false, true, false, false), false)]
#[case::session_trust_env_beats_opt_out(proxy_flags(true, true, false, false), true)]
#[case::http2_uses_httpx(proxy_flags(false, true, false, true), true)]
#[case::aiohttp_disabled(proxy_flags(false, true, true, false), true)]
fn environment_proxies_apply_unless_the_aiohttp_transport_opts_out(
#[case] layer: HttpSettingsLayer,
#[case] expected: bool,
) {
assert_eq!(HttpSettings::from_layers([layer]).trust_proxy_env, expected);
}
#[test]
fn a_proxy_opt_out_in_one_source_still_yields_to_trust_env_from_another() {
let environment =
HttpSettingsLayer::from_environment(&env_of(&[("DISABLE_AIOHTTP_TRUST_ENV", "true")]));
let configured = HttpSettingsLayer {
aiohttp_trust_env: Some(true),
..HttpSettingsLayer::default()
};
assert!(!HttpSettings::from_layers([environment.clone()]).trust_proxy_env);
assert!(HttpSettings::from_layers([environment, configured]).trust_proxy_env);
}
#[test]
fn missing_files_fall_back_to_default_verification() {
let settings = HttpSettings {
ssl_verify: Some(SslVerify::CaBundle("/absent/roots.pem".into())),
ssl_cert_file: Some("/absent/env.pem".into()),
..HttpSettings::default()
}
.without_missing_files(&|_| false);
assert_eq!(settings.ssl_verify, Some(SslVerify::Enabled));
assert_eq!(settings.ssl_cert_file, None);
}
#[test]
fn existing_files_are_kept() {
let settings = HttpSettings {
ssl_verify: Some(SslVerify::CaBundle("/present/roots.pem".into())),
ssl_cert_file: Some("/present/env.pem".into()),
..HttpSettings::default()
};
assert_eq!(settings.clone().without_missing_files(&|_| true), settings);
}
#[rstest]
#[case("true", Some(true))]
#[case("True", Some(true))]
#[case("false", None)]
#[case("1", None)]
fn boolean_switches_only_turn_on_for_true(
#[case] value: &'static str,
#[case] expected: Option<bool>,
) {
let env = move |name: &str| match name {
"LITELLM_HTTP2"
| "AIOHTTP_TRUST_ENV"
| "DISABLE_AIOHTTP_TRANSPORT"
| "DISABLE_AIOHTTP_TRUST_ENV" => Some(value.to_string()),
_ => None,
};
let layer = HttpSettingsLayer::from_environment(&env);
assert_eq!(layer.http2, expected);
assert_eq!(layer.aiohttp_trust_env, expected);
assert_eq!(layer.disable_aiohttp_transport, expected);
assert_eq!(layer.disable_aiohttp_trust_env, expected);
}
}

411
litellm-rust/crates/http/src/tls.rs Normal file
View file

@ -0,0 +1,411 @@
use std::{fmt, path::Path, str::FromStr, sync::Arc};
use rustls::{
CipherSuite, ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme,
client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
crypto::{CryptoProvider, SupportedKxGroup, ring},
pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime, pem::PemObject},
};
use crate::{
config::{HttpClientConfig, Verify},
error::Error,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum KeyExchangeGroup {
X25519,
Secp256r1,
Secp384r1,
}
impl FromStr for KeyExchangeGroup {
type Err = Unsupported;
fn from_str(name: &str) -> Result<Self, Self::Err> {
match name.trim().to_ascii_lowercase().as_str() {
"x25519" => Ok(Self::X25519),
"prime256v1" | "secp256r1" | "p-256" => Ok(Self::Secp256r1),
"secp384r1" | "p-384" => Ok(Self::Secp384r1),
_ => Err(Unsupported::EcdhCurve(name.to_owned())),
}
}
}
impl KeyExchangeGroup {
fn supported(self) -> &'static dyn SupportedKxGroup {
match self {
Self::X25519 => ring::kx_group::X25519,
Self::Secp256r1 => ring::kx_group::SECP256R1,
Self::Secp384r1 => ring::kx_group::SECP384R1,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum Tls12CipherSuite {
EcdheEcdsaAes128Gcm,
EcdheEcdsaAes256Gcm,
EcdheEcdsaChacha20,
EcdheRsaAes128Gcm,
EcdheRsaAes256Gcm,
EcdheRsaChacha20,
}
impl FromStr for Tls12CipherSuite {
type Err = Unsupported;
fn from_str(name: &str) -> Result<Self, Self::Err> {
match name {
"ECDHE-ECDSA-AES128-GCM-SHA256" => Ok(Self::EcdheEcdsaAes128Gcm),
"ECDHE-ECDSA-AES256-GCM-SHA384" => Ok(Self::EcdheEcdsaAes256Gcm),
"ECDHE-ECDSA-CHACHA20-POLY1305" => Ok(Self::EcdheEcdsaChacha20),
"ECDHE-RSA-AES128-GCM-SHA256" => Ok(Self::EcdheRsaAes128Gcm),
"ECDHE-RSA-AES256-GCM-SHA384" => Ok(Self::EcdheRsaAes256Gcm),
"ECDHE-RSA-CHACHA20-POLY1305" => Ok(Self::EcdheRsaChacha20),
_ => Err(Unsupported::CipherToken(name.to_owned())),
}
}
}
impl Tls12CipherSuite {
fn suite(self) -> CipherSuite {
match self {
Self::EcdheEcdsaAes128Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
Self::EcdheEcdsaAes256Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
Self::EcdheEcdsaChacha20 => CipherSuite::TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,
Self::EcdheRsaAes128Gcm => CipherSuite::TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
Self::EcdheRsaAes256Gcm => CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
Self::EcdheRsaChacha20 => CipherSuite::TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, thiserror::Error)]
pub enum Unsupported {
#[error(
"ssl_ecdh_curve {0:?} is not supported: rustls with ring only offers X25519, prime256v1 and secp384r1, so the default key exchange groups are used"
)]
EcdhCurve(String),
#[error(
"ssl_security_level {0:?} is not supported: rustls has one fixed security level, comparable to OpenSSL level 2, so legacy servers that need a lower level cannot be reached"
)]
SecurityLevel(String),
#[error(
"ssl_security_level entry {0:?} is not supported: rustls only offers ECDHE AEAD cipher suites, so the entry is ignored"
)]
CipherToken(String),
}
#[derive(Default)]
pub(crate) struct CipherSelection {
pub(crate) tls12_cipher_suites: Option<Vec<Tls12CipherSuite>>,
pub(crate) unsupported: Vec<Unsupported>,
}
enum CipherToken {
Suite(Tls12CipherSuite),
EverySuite,
Ordering,
Unsupported(Unsupported),
}
impl From<&str> for CipherToken {
fn from(token: &str) -> Self {
match token {
"DEFAULT" | "ALL" | "HIGH" => Self::EverySuite,
"@STRENGTH" | "@SECLEVEL=2" => Self::Ordering,
level if level.starts_with("@SECLEVEL=") => {
Self::Unsupported(Unsupported::SecurityLevel(level.to_owned()))
}
name => name.parse().map_or_else(Self::Unsupported, Self::Suite),
}
}
}
impl From<&str> for CipherSelection {
fn from(value: &str) -> Self {
let tokens: Vec<CipherToken> = tokenize(value)
.iter()
.map(|token| CipherToken::from(token.as_str()))
.collect();
let every_suite = tokens
.iter()
.any(|token| matches!(token, CipherToken::EverySuite));
let mut suites: Vec<Tls12CipherSuite> = tokens
.iter()
.filter_map(|token| match token {
CipherToken::Suite(suite) => Some(*suite),
_ => None,
})
.collect();
suites.sort_unstable();
suites.dedup();
CipherSelection {
tls12_cipher_suites: (!every_suite && !suites.is_empty()).then_some(suites),
unsupported: tokens
.into_iter()
.filter_map(|token| match token {
CipherToken::Unsupported(unsupported) => Some(unsupported),
_ => None,
})
.collect(),
}
}
}
fn tokenize(value: &str) -> Vec<String> {
value
.split([':', ',', ' '])
.flat_map(|entry| match entry.split_once('@') {
Some((name, command)) => vec![name.to_owned(), format!("@{command}")],
None => vec![entry.to_owned()],
})
.filter(|token| !token.is_empty())
.collect()
}
impl TryFrom<&HttpClientConfig> for ClientConfig {
type Error = Error;
fn try_from(config: &HttpClientConfig) -> Result<Self, Self::Error> {
let base = ring::default_provider();
let provider = Arc::new(CryptoProvider {
kx_groups: config
.key_exchange_group
.map_or_else(|| base.kx_groups.clone(), |group| vec![group.supported()]),
cipher_suites: base
.cipher_suites
.iter()
.copied()
.filter(|suite| {
suite.tls13().is_some()
|| config.tls12_cipher_suites.as_ref().is_none_or(|allowed| {
allowed.iter().any(|a| a.suite() == suite.suite())
})
})
.collect(),
..base
});
let builder = ClientConfig::builder_with_provider(Arc::clone(&provider))
.with_safe_default_protocol_versions()
.map_err(|error| Error::Client(error.to_string()))?;
let verified = match &config.verify {
Verify::Disabled => builder
.dangerous()
.with_custom_certificate_verifier(Arc::new(NoVerification(provider))),
Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore {
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
}),
Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?),
};
let mut tls = match &config.client_certificate {
None => verified.with_no_client_auth(),
Some(path) => {
let (chain, key) = identity(path)?;
verified
.with_client_auth_cert(chain, key)
.map_err(|error| invalid_pem(path, error))?
}
};
tls.alpn_protocols = if config.http2 {
vec![b"h2".to_vec(), b"http/1.1".to_vec()]
} else {
vec![b"http/1.1".to_vec()]
};
Ok(tls)
}
}
fn bundle_roots(path: &Path) -> Result<RootCertStore, Error> {
let certificates = certificates(path)?;
if certificates.is_empty() {
return Err(invalid_pem(path, "no certificates found"));
}
let mut store = RootCertStore::empty();
for certificate in certificates {
store
.add(certificate)
.map_err(|error| invalid_pem(path, error))?;
}
Ok(store)
}
fn identity(path: &Path) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>), Error> {
let chain = certificates(path)?;
if chain.is_empty() {
return Err(invalid_pem(path, "no certificates found"));
}
let key =
PrivateKeyDer::from_pem_slice(&read(path)?).map_err(|error| invalid_pem(path, error))?;
Ok((chain, key))
}
fn certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, Error> {
CertificateDer::pem_slice_iter(&read(path)?)
.collect::<Result<_, _>>()
.map_err(|error| invalid_pem(path, error))
}
fn read(path: &Path) -> Result<Vec<u8>, Error> {
std::fs::read(path).map_err(|error| Error::Read {
path: path.to_path_buf(),
message: error.to_string(),
})
}
fn invalid_pem(path: &Path, message: impl fmt::Display) -> Error {
Error::InvalidPem {
path: path.to_path_buf(),
message: message.to_string(),
}
}
#[derive(Debug)]
struct NoVerification(Arc<CryptoProvider>);
impl ServerCertVerifier for NoVerification {
fn verify_server_cert(
&self,
_end_entity: &CertificateDer<'_>,
_intermediates: &[CertificateDer<'_>],
_server_name: &ServerName<'_>,
_ocsp_response: &[u8],
_now: UnixTime,
) -> Result<ServerCertVerified, rustls::Error> {
Ok(ServerCertVerified::assertion())
}
fn verify_tls12_signature(
&self,
_message: &[u8],
_cert: &CertificateDer<'_>,
_dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
Ok(HandshakeSignatureValid::assertion())
}
fn verify_tls13_signature(
&self,
_message: &[u8],
_cert: &CertificateDer<'_>,
_dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
Ok(HandshakeSignatureValid::assertion())
}
fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
self.0.signature_verification_algorithms.supported_schemes()
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use rustls::NamedGroup;
use super::*;
use crate::{HttpSettings, Resolution};
fn config(settings: HttpSettings) -> HttpClientConfig {
Resolution::from(&settings).config
}
fn offered_groups(tls: &ClientConfig) -> Vec<NamedGroup> {
tls.crypto_provider()
.kx_groups
.iter()
.map(|group| group.name())
.collect()
}
fn offered_tls12_suites(tls: &ClientConfig) -> Vec<CipherSuite> {
tls.crypto_provider()
.cipher_suites
.iter()
.filter(|suite| suite.tls13().is_none())
.map(|suite| suite.suite())
.collect()
}
#[rstest]
#[case("X25519", NamedGroup::X25519)]
#[case("prime256v1", NamedGroup::secp256r1)]
#[case("secp384r1", NamedGroup::secp384r1)]
fn ecdh_curve_is_the_only_key_exchange_group_offered(
#[case] curve: &str,
#[case] expected: NamedGroup,
) {
let tls = ClientConfig::try_from(&config(HttpSettings {
ssl_ecdh_curve: Some(curve.into()),
..HttpSettings::default()
}))
.unwrap();
assert_eq!(offered_groups(&tls), [expected]);
}
#[test]
fn default_settings_offer_every_group_and_suite_of_the_provider() {
let tls = ClientConfig::try_from(&config(HttpSettings::default())).unwrap();
let provider = ring::default_provider();
assert_eq!(offered_groups(&tls).len(), provider.kx_groups.len());
assert_eq!(
tls.crypto_provider().cipher_suites.len(),
provider.cipher_suites.len()
);
}
#[test]
fn named_suites_are_the_only_tls12_suites_offered_and_tls13_stays() {
let tls = ClientConfig::try_from(&config(HttpSettings {
ssl_security_level: Some("ECDHE-RSA-AES256-GCM-SHA384".into()),
..HttpSettings::default()
}))
.unwrap();
assert_eq!(
offered_tls12_suites(&tls),
[CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384]
);
assert!(
tls.crypto_provider()
.cipher_suites
.iter()
.any(|suite| suite.tls13().is_some())
);
}
#[rstest]
#[case(true, &[b"h2".as_slice(), b"http/1.1".as_slice()])]
#[case(false, &[b"http/1.1".as_slice()])]
fn alpn_offers_h2_only_when_http2_is_on(#[case] http2: bool, #[case] expected: &[&[u8]]) {
let tls = ClientConfig::try_from(&config(HttpSettings {
http2,
..HttpSettings::default()
}))
.unwrap();
assert_eq!(tls.alpn_protocols, expected);
}
#[test]
fn client_certificate_without_a_private_key_is_an_invalid_pem_error() {
let path = std::env::temp_dir().join(format!(
"litellm-http-cert-without-key-{}.pem",
std::process::id()
));
std::fs::write(
&path,
b"-----BEGIN CERTIFICATE-----\nAA==\n-----END CERTIFICATE-----\n",
)
.unwrap();
let result = ClientConfig::try_from(&HttpClientConfig {
client_certificate: Some(path.clone()),
..config(HttpSettings::default())
})
.map(drop);
std::fs::remove_file(&path).unwrap();
assert!(matches!(
result,
Err(Error::InvalidPem { path: reported, .. }) if reported == path
));
}
}

1
litellm-rust/crates/llms/Cargo.toml
View file

@ -17,6 +17,7 @@ litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
litellm-host.workspace = true
litellm-framing.workspace = true
litellm-http.workspace = true
base64.workspace = true
bytes.workspace = true
data-url = "0.3.2"

1
litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs
View file

@ -21,7 +21,6 @@ use crate::{
pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024;
pub const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
pub const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
pub const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024;
pub const OCR_MAX_FETCH_REDIRECTS: usize = 10;

55
litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs
View file

@ -1,9 +1,8 @@
use std::{sync::OnceLock, time::Duration};
use bytes::{Bytes, BytesMut};
use futures_util::future::BoxFuture;
use litellm_auth_gcp::VertexAuth;
use litellm_host::event::WireRequest;
use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool};
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
@ -11,14 +10,13 @@ use crate::{
base_llm::ocr::{
error::Error,
transformation::{
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_CONNECT_TIMEOUT_SECS,
OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value,
decode_response,
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext,
PreparedOcrRequest, decode_request_value, decode_response,
},
},
custom_httpx::{
http_handler::{HeaderPolicy, execute_http_request, with_headers},
media::MediaFetcher,
media::{MediaFetcher, UrlPolicy},
transport,
},
};
@ -40,30 +38,20 @@ pub struct OcrClient {
}
impl OcrClient {
pub fn new(provider_http: reqwest::Client) -> Result<Self, transport::Error> {
let document_fetcher = MediaFetcher::new().map_err(transport::Error::from)?;
pub fn new(
pool: &HttpClientPool,
config: &HttpClientConfig,
url_policy: UrlPolicy,
vertex_auth: VertexAuth,
) -> Result<Self, litellm_http::Error> {
Ok(Self {
provider_http,
polling_http: no_redirect_http()?,
document_fetcher,
vertex_auth: VertexAuth::default(),
provider_http: pool.client(config, ClientVariant::Provider)?,
polling_http: pool.client(config, ClientVariant::NoRedirect)?,
document_fetcher: MediaFetcher::new(pool, config, url_policy)?,
vertex_auth,
})
}
pub fn shared() -> Result<Self, Error> {
static CLIENT: OnceLock<Result<OcrClient, transport::Error>> = OnceLock::new();
let client = CLIENT
.get_or_init(|| {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
.build()
.map_err(transport::Error::from)
.and_then(OcrClient::new)
})
.clone()?;
Ok(client)
}
pub fn provider_http(&self) -> &reqwest::Client {
&self.provider_http
}
@ -84,21 +72,16 @@ impl OcrClient {
pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self {
Self {
provider_http,
polling_http: no_redirect_http().expect("test polling client builds"),
polling_http: reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test polling client builds"),
document_fetcher: MediaFetcher::for_test(document_http),
vertex_auth: VertexAuth::default(),
}
}
}
fn no_redirect_http() -> Result<reqwest::Client, transport::Error> {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(transport::Error::from)
}
/// Rust counterpart of `BaseLLMHTTPHandler.async_ocr`: prepare the provider request,
/// send it, and hand the response to the config for normalization.
pub async fn ocr<C: BaseOcrConfig>(
@ -296,6 +279,8 @@ pub fn body_document(body: &Value) -> Result<OcrDocument, Error> {
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
#[tokio::test]

232
litellm-rust/crates/llms/src/custom_httpx/media.rs
View file

@ -7,13 +7,12 @@ use std::{
time::Duration,
};
use litellm_http::{ClientVariant, EnvironmentProxies, HttpClientConfig, HttpClientPool};
use reqwest::{
Url,
dns::{Addrs, Name, Resolve, Resolving},
};
const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("media URL rejected by network policy")]
@ -36,10 +35,45 @@ pub enum Error {
Transport(#[from] crate::custom_httpx::transport::Error),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UrlPolicy {
pub validate: bool,
pub allowed_hosts: Vec<String>,
}
impl Default for UrlPolicy {
fn default() -> Self {
Self {
validate: true,
allowed_hosts: Vec::new(),
}
}
}
impl UrlPolicy {
fn allows(&self, host: &str, port: u16) -> bool {
let host = normalize_host(host);
let with_port = format!("{host}:{port}");
self.allowed_hosts
.iter()
.map(|entry| normalize_host(entry))
.any(|entry| entry == host || entry == with_port)
}
}
fn normalize_host(host: &str) -> String {
host.to_ascii_lowercase().trim_end_matches('.').to_owned()
}
type ProxyMatch = Arc<dyn Fn(&Url) -> bool + Send + Sync>;
#[derive(Clone)]
pub struct MediaFetcher {
client: reqwest::Client,
pinned: reqwest::Client,
unpinned: reqwest::Client,
uses_proxy: ProxyMatch,
address_resolver: Arc<dyn AddressResolver>,
url_policy: UrlPolicy,
allow_private_network: bool,
}
@ -63,26 +97,39 @@ pub struct DownloadedMedia {
}
impl MediaFetcher {
pub fn new() -> Result<Self, reqwest::Error> {
Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver))
pub fn new(
pool: &HttpClientPool,
config: &HttpClientConfig,
url_policy: UrlPolicy,
) -> Result<Self, litellm_http::Error> {
let uses_proxy: ProxyMatch = if config.trust_proxy_env {
let proxies = EnvironmentProxies::from_environment();
Arc::new(move |url| proxies.apply_to(url))
} else {
Arc::new(|_| false)
};
Self::with_resolution(
pool,
config,
url_policy,
Arc::new(SystemAddressResolver),
uses_proxy,
)
}
fn with_resolvers<R>(
transport_resolver: Arc<R>,
fn with_resolution(
pool: &HttpClientPool,
config: &HttpClientConfig,
url_policy: UrlPolicy,
address_resolver: Arc<dyn AddressResolver>,
) -> Result<Self, reqwest::Error>
where
R: Resolve + 'static,
{
let client = reqwest::Client::builder()
.connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.no_proxy()
.dns_resolver(transport_resolver)
.build()?;
uses_proxy: ProxyMatch,
) -> Result<Self, litellm_http::Error> {
Ok(Self {
client,
pinned: pool.client(config, ClientVariant::Media)?,
unpinned: pool.client(config, ClientVariant::UnpinnedMedia)?,
uses_proxy,
address_resolver,
url_policy,
allow_private_network: false,
})
}
@ -90,8 +137,11 @@ impl MediaFetcher {
#[cfg(any(test, feature = "test-support"))]
pub fn for_test(client: reqwest::Client) -> Self {
Self {
client,
pinned: client.clone(),
unpinned: client,
uses_proxy: Arc::new(|_| false),
address_resolver: Arc::new(AllowPrivateResolver),
url_policy: UrlPolicy::default(),
allow_private_network: true,
}
}
@ -112,9 +162,9 @@ impl MediaFetcher {
) -> Result<DownloadedMedia, Error> {
let mut redirects_followed = 0;
loop {
self.validate_url(&url).await?;
let mut response = self
.client
.client_for(&url)
.await?
.get(url.clone())
.send()
.await
@ -161,7 +211,10 @@ impl MediaFetcher {
}
}
async fn validate_url(&self, url: &Url) -> Result<(), Error> {
async fn client_for(&self, url: &Url) -> Result<&reqwest::Client, Error> {
if !self.url_policy.validate {
return Ok(&self.unpinned);
}
if !matches!(url.scheme(), "http" | "https")
|| !url.username().is_empty()
|| url.password().is_some()
@ -170,12 +223,28 @@ impl MediaFetcher {
}
let host = url.host_str().ok_or(Error::BlockedUrl)?;
if self.allow_private_network {
return Ok(());
}
if let Ok(ip) = host.parse::<IpAddr>() {
return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl);
return Ok(&self.pinned);
}
let port = url.port_or_known_default().ok_or(Error::BlockedUrl)?;
if self.url_policy.allows(host, port) {
return Ok(&self.unpinned);
}
self.validate_host(host, port).await?;
Ok(if (self.uses_proxy)(url) {
&self.unpinned
} else {
&self.pinned
})
}
async fn validate_host(&self, host: &str, port: u16) -> Result<(), Error> {
if let Ok(ip) = host
.trim_start_matches('[')
.trim_end_matches(']')
.parse::<IpAddr>()
{
return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl);
}
let addresses = self
.address_resolver
.resolve(host, port)
@ -236,7 +305,7 @@ fn is_blocked_ip(ip: IpAddr) -> bool {
}
#[derive(Default)]
struct PublicDnsResolver;
pub struct PublicDnsResolver;
struct SystemAddressResolver;
@ -281,6 +350,7 @@ impl Resolve for PublicDnsResolver {
mod tests {
use std::collections::HashSet;
use litellm_http::{HttpSettings, Resolution};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
@ -364,13 +434,35 @@ mod tests {
address: SocketAddr,
blocked_hosts: HashSet<&'static str>,
) -> MediaFetcher {
MediaFetcher::with_resolvers(
Arc::new(LoopbackDnsResolver(address)),
fetcher(address, blocked_hosts, UrlPolicy::default(), false)
}
fn fetcher(
pinned_address: SocketAddr,
blocked_hosts: HashSet<&'static str>,
url_policy: UrlPolicy,
uses_proxy: bool,
) -> MediaFetcher {
let direct = HttpClientConfig {
trust_proxy_env: false,
..Resolution::from(&HttpSettings::default()).config
};
MediaFetcher::with_resolution(
&HttpClientPool::new(Arc::new(LoopbackDnsResolver(pinned_address))),
&direct,
url_policy,
Arc::new(TestAddressResolver { blocked_hosts }),
Arc::new(move |_| uses_proxy),
)
.expect("test fetcher builds")
}
const UNROUTABLE: SocketAddr =
SocketAddr::new(IpAddr::V4(std::net::Ipv4Addr::new(192, 0, 2, 1)), 9);
const OK_RESPONSE: &[u8] =
b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok";
fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy {
DownloadPolicy {
timeout: Duration::from_secs(1),
@ -542,12 +634,90 @@ mod tests {
#[tokio::test]
async fn rejects_url_credentials_before_network_access() {
let fetcher = MediaFetcher::new().expect("media fetcher builds");
let fetcher = MediaFetcher::new(
&HttpClientPool::new(Arc::new(PublicDnsResolver)),
&Resolution::from(&HttpSettings::default()).config,
UrlPolicy::default(),
)
.expect("media fetcher builds");
let url =
Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses");
assert!(matches!(
fetcher.validate_url(&url).await,
fetcher.fetch(url, policy(1, 0)).await,
Err(Error::BlockedUrl)
));
}
#[tokio::test]
async fn allowlisted_private_host_is_fetched_without_the_pinned_resolver() {
let (url, server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
let port = url.port().expect("test URL has a port");
let allowed = UrlPolicy {
validate: true,
allowed_hosts: vec![format!("LOCALHOST:{port}")],
};
let media = fetcher(UNROUTABLE, HashSet::from(["localhost"]), allowed, false)
.fetch(url, policy(2, 0))
.await
.expect("allowlisted host downloads");
server.await.expect("server completes");
assert_eq!(media.bytes, b"ok");
}
#[tokio::test]
async fn allowlist_entry_for_another_port_does_not_open_the_host() {
let (url, _server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
let other_port = UrlPolicy {
validate: true,
allowed_hosts: vec!["localhost:1".into()],
};
let result = fetcher(UNROUTABLE, HashSet::from(["localhost"]), other_port, false)
.fetch(url, policy(2, 0))
.await;
assert!(matches!(result, Err(Error::BlockedUrl)));
}
#[tokio::test]
async fn validation_off_fetches_private_hosts_and_follows_redirects() {
let (url, server, _) = serve_named(
"localhost",
vec![
b"HTTP/1.1 302 Found\r\nLocation: /moved\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
OK_RESPONSE,
],
)
.await;
let off = UrlPolicy {
validate: false,
allowed_hosts: Vec::new(),
};
let media = fetcher(UNROUTABLE, HashSet::from(["localhost"]), off, false)
.fetch(url, policy(2, 1))
.await
.expect("unvalidated download succeeds");
let requests = server.await.expect("server completes");
assert_eq!(media.bytes, b"ok");
assert!(requests[1].starts_with("GET /moved "));
}
#[tokio::test]
async fn proxied_urls_skip_the_pinned_resolver_but_keep_the_address_check() {
let (url, server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
let media = fetcher(UNROUTABLE, HashSet::new(), UrlPolicy::default(), true)
.fetch(url.clone(), policy(2, 0))
.await
.expect("public host behind a proxy downloads");
server.await.expect("server completes");
assert_eq!(media.bytes, b"ok");
let blocked = fetcher(
UNROUTABLE,
HashSet::from(["localhost"]),
UrlPolicy::default(),
true,
)
.fetch(url, policy(2, 0))
.await;
assert!(matches!(blocked, Err(Error::BlockedUrl)));
}
}

38
litellm-rust/crates/llms/src/custom_httpx/transport.rs
View file

@ -11,7 +11,7 @@ pub enum Error {
impl Error {
pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self {
let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder());
let message = error.without_url().to_string();
let message = describe(error);
if before_dispatch {
Self::Connect(message)
} else {
@ -22,10 +22,18 @@ impl Error {
impl From<reqwest::Error> for Error {
fn from(error: reqwest::Error) -> Self {
Self::Network(error.without_url().to_string())
Self::Network(describe(error))
}
}
fn describe(error: reqwest::Error) -> String {
let error = error.without_url();
std::iter::successors(std::error::Error::source(&error), |cause| cause.source())
.fold(error.to_string(), |message, cause| {
format!("{message}: {cause}")
})
}
#[cfg(test)]
mod tests {
#[tokio::test]
@ -47,6 +55,32 @@ mod tests {
assert!(!error.to_string().contains("private"));
}
fn root_cause(error: &dyn std::error::Error) -> Option<String> {
match error.source() {
Some(cause) => root_cause(cause).or_else(|| Some(cause.to_string())),
None => None,
}
}
#[tokio::test]
async fn network_error_message_names_the_underlying_cause() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
let address = listener.local_addr().expect("address");
drop(listener);
let error = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}/private?api_key=secret"))
.send()
.await
.expect_err("nothing listens on the port");
let root_cause = root_cause(&error).expect("reqwest reports a cause");
let message = crate::custom_httpx::transport::Error::from(error).to_string();
assert!(message.contains(&root_cause), "{message}");
assert!(!message.contains("secret"));
}
#[tokio::test]
async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() {
use std::time::Duration;

44
litellm-rust/crates/python-bridge/AGENTS.md
View file

@ -1,4 +1,4 @@
- Target invariants, not completion claims; these supersede older conflicting bridge guidance
- Target invariants, not completion claims; these supersede the crate guidance below where they conflict
- Keep this crate the product-specific PyO3 consumer of `litellm-host-python`
- Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract
- Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy
@ -34,3 +34,45 @@
- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [exception transfer](https://docs.rs/pyo3/0.29.2/pyo3/struct.PyErr.html#method.into_value), [re-entry](https://pyo3.rs/v0.29.2/class/call.html)
- [GIL policy](https://pyo3.rs/v0.29.2/free-threading.html), [experimental async limits](https://pyo3.rs/v0.29.2/async-await.html), [task conversion](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/fn.into_future_with_locals.html), [native cancellation/delivery](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/tokio/fn.future_into_py.html)
- [performance](https://pyo3.rs/v0.29.2/performance.html), [PyBackedBytes](https://docs.rs/pyo3/0.29.2/pyo3/pybacked/struct.PyBackedBytes.html), [typing](https://pyo3.rs/v0.29.2/python-typing-hints.html)
Rules for `litellm-rust/crates/python-bridge`.
## Responsibility
`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms.
Keep this crate thin. It exposes LiteLLM Rust APIs, assembles domain requests,
maps domain errors to Python exceptions, and delegates generic conversion and
GIL handling to `litellm-host-python`.
## Bridge Shape
- Prefer one stable method per top-level LiteLLM route, for example
`messages(...)`, calling the matching `litellm-core` entrypoint.
- Do not add one exported PyO3 function per provider helper unless there is a
measured reason.
- Provider dispatch belongs in the `litellm-core` route module (e.g.
`litellm_core::messages`), not in this PyO3 crate.
- Python owns rollout state and fallback. Rust should return errors; Python
decides whether to raise or fall back. For a rust-only provider/route (no
Python reference), the Python side is a thin dispatch that calls Rust and
raises when the bridge is unavailable, with no fallback.
- Keep the Python interface minimal (well under 100 lines per route): it only
marshals inputs and calls Rust. Do not add per-route feature flags, and do
not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch
class under `litellm/llms/<provider>/<route>/`.
## Data Handling
- OCR payloads can contain personal data and large base64 images. Do not log
payloads or provider responses.
- Avoid copying large payloads more than needed. The current JSON round-trip is
acceptable for the first scaffold, but future performance work should evaluate
direct PyO3 conversion before expanding Rust coverage to image-heavy paths.
- Do not expose raw Rust errors that include document contents or upstream
bodies.
## Tests
- `cargo test --workspace` must compile this crate.
- Python tests must cover bridge disabled, bridge enabled, and module-missing
fallback behavior for every exposed route.

43
litellm-rust/crates/python-bridge/CLAUDE.md
View file

@ -1,43 +0,0 @@
# CLAUDE.md
Rules for `litellm-rust/crates/python-bridge`.
## Responsibility
`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms.
Keep this crate thin. It exposes LiteLLM Rust APIs, assembles domain requests,
maps domain errors to Python exceptions, and delegates generic conversion and
GIL handling to `litellm-host-python`.
## Bridge Shape
- Prefer one stable method per top-level LiteLLM route, for example
`messages(...)`, calling the matching `litellm-core` entrypoint.
- Do not add one exported PyO3 function per provider helper unless there is a
measured reason.
- Provider dispatch belongs in the `litellm-core` route module (e.g.
`litellm_core::messages`), not in this PyO3 crate.
- Python owns rollout state and fallback. Rust should return errors; Python
decides whether to raise or fall back. For a rust-only provider/route (no
Python reference), the Python side is a thin dispatch that calls Rust and
raises when the bridge is unavailable, with no fallback.
- Keep the Python interface minimal (well under 100 lines per route): it only
marshals inputs and calls Rust. Do not add per-route feature flags, and do
not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch
class under `litellm/llms/<provider>/<route>/`.
## Data Handling
- OCR payloads can contain personal data and large base64 images. Do not log
payloads or provider responses.
- Avoid copying large payloads more than needed. The current JSON round-trip is
acceptable for the first scaffold, but future performance work should evaluate
direct PyO3 conversion before expanding Rust coverage to image-heavy paths.
- Do not expose raw Rust errors that include document contents or upstream
bodies.
## Tests
- `cargo test --workspace` must compile this crate.
- Python tests must cover bridge disabled, bridge enabled, and module-missing
fallback behavior for every exposed route.

2
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -20,6 +20,8 @@ bytes.workspace = true
litellm-auth.workspace = true
litellm-callbacks-legacy.workspace = true
litellm-core.workspace = true
litellm-auth-gcp.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
litellm-types.workspace = true
litellm-host-python.workspace = true

18
litellm-rust/crates/python-bridge/python_settings.json Normal file
View file

@ -0,0 +1,18 @@
{
"http_settings": [
"ssl_verify",
"ssl_certificate",
"ssl_security_level",
"ssl_ecdh_curve",
"force_ipv4",
"http2",
"aiohttp_trust_env",
"disable_aiohttp_trust_env",
"disable_aiohttp_transport",
"user_agent"
],
"url_policy": [
"user_url_validation",
"user_url_allowed_hosts"
]
}

354
litellm-rust/crates/python-bridge/src/http.rs Normal file
View file

@ -0,0 +1,354 @@
use std::{
collections::HashSet,
path::{Path, PathBuf},
sync::{Arc, LazyLock, Mutex, PoisonError},
};
use litellm_http::{
HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify,
Unsupported,
};
use litellm_llms::custom_httpx::media::{PublicDnsResolver, UrlPolicy};
use pyo3::{prelude::*, types::PyDict};
use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings};
static POOL: LazyLock<HttpClientPool> =
LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver)));
static REPORTED_UNSUPPORTED: LazyLock<Mutex<HashSet<Unsupported>>> = LazyLock::new(Mutex::default);
pub(crate) fn pool() -> &'static HttpClientPool {
&POOL
}
pub(crate) fn call_config(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<HttpClientConfig> {
let settings = HttpSettings::from_layers([
for_call(call_ssl_verify(kwargs)?, asynchronous),
HttpSettingsLayer::from_environment(&|name| std::env::var(name).ok()),
configured(&PythonSettings::Http.read(py)?)?,
])
.without_missing_files(&|path: &Path| path.exists());
let resolution = Resolution::from(&settings);
for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) {
PythonSettings::warn(py, &unsupported.to_string())?;
}
Ok(resolution.config)
}
fn unreported(
reported: &Mutex<HashSet<Unsupported>>,
unsupported: Vec<Unsupported>,
) -> Vec<Unsupported> {
let mut reported = reported.lock().unwrap_or_else(PoisonError::into_inner);
unsupported
.into_iter()
.filter(|unsupported| reported.insert(unsupported.clone()))
.collect()
}
pub(crate) fn url_policy(py: Python<'_>) -> PyResult<UrlPolicy> {
let policy: PythonUrlPolicy =
PythonSettings::UrlPolicy
.read(py)?
.extract()
.map_err(|error: PyErr| {
RustBridgeDeclined::new_err(format!(
"litellm URL policy cannot be used by the Rust route: {error}"
))
})?;
Ok(UrlPolicy {
validate: policy.user_url_validation,
allowed_hosts: policy.user_url_allowed_hosts,
})
}
fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult<Option<SslVerify>> {
Ok(kwargs
.get_item("ssl_verify")?
.and_then(|value| ssl_verify(&value)))
}
fn for_call(call_ssl_verify: Option<SslVerify>, asynchronous: bool) -> HttpSettingsLayer {
HttpSettingsLayer {
ssl_verify: call_ssl_verify,
disable_aiohttp_transport: (!asynchronous).then_some(true),
..HttpSettingsLayer::default()
}
}
#[derive(FromPyObject)]
struct PythonUrlPolicy {
user_url_validation: bool,
user_url_allowed_hosts: Vec<String>,
}
#[derive(FromPyObject)]
struct PythonHttpSettings<'py> {
ssl_verify: Bound<'py, PyAny>,
ssl_certificate: Option<String>,
ssl_security_level: Option<String>,
ssl_ecdh_curve: Option<String>,
force_ipv4: bool,
http2: bool,
aiohttp_trust_env: bool,
disable_aiohttp_trust_env: bool,
disable_aiohttp_transport: bool,
user_agent: String,
}
fn configured(value: &Bound<'_, PyAny>) -> PyResult<HttpSettingsLayer> {
let python: PythonHttpSettings = value.extract().map_err(|error: PyErr| {
RustBridgeDeclined::new_err(format!(
"litellm HTTP settings cannot be used by the Rust route: {error}"
))
})?;
Ok(HttpSettingsLayer {
ssl_verify: ssl_verify(&python.ssl_verify),
ssl_certificate: python.ssl_certificate.map(PathBuf::from),
ssl_security_level: python.ssl_security_level,
ssl_ecdh_curve: python.ssl_ecdh_curve,
force_ipv4: Some(python.force_ipv4),
http2: Some(python.http2),
aiohttp_trust_env: Some(python.aiohttp_trust_env),
disable_aiohttp_trust_env: Some(python.disable_aiohttp_trust_env),
disable_aiohttp_transport: Some(python.disable_aiohttp_transport),
user_agent: Some(python.user_agent),
..HttpSettingsLayer::default()
})
}
fn ssl_verify(value: &Bound<'_, PyAny>) -> Option<SslVerify> {
if let Ok(enabled) = value.extract::<bool>() {
return Some(if enabled {
SslVerify::Enabled
} else {
SslVerify::Disabled
});
}
value
.extract::<String>()
.ok()
.map(|path| SslVerify::parse(&path))
}
#[cfg(test)]
mod tests {
use litellm_http::Verify;
use rstest::rstest;
use super::*;
use crate::python_settings::CONTRACT;
fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> {
let source = format!(
"
import json
import types
defaults = dict(
ssl_verify=True,
ssl_certificate=None,
ssl_security_level=None,
ssl_ecdh_curve=None,
force_ipv4=False,
http2=False,
aiohttp_trust_env=False,
disable_aiohttp_trust_env=False,
disable_aiohttp_transport=False,
user_agent='litellm/test',
)
defaults.update(dict({overrides}))
settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']}})
"
);
let locals = PyDict::new(py);
locals.set_item("contract", CONTRACT).unwrap();
let source = std::ffi::CString::new(source).unwrap();
py.run(&source, Some(&locals), Some(&locals)).unwrap();
locals.get_item("settings").unwrap().unwrap()
}
#[test]
fn default_python_settings_resolve_to_default_settings_with_verification_on() {
Python::initialize();
Python::attach(|py| {
let layer = configured(&python_settings(py, "")).unwrap();
assert_eq!(
HttpSettings::from_layers([layer]),
HttpSettings {
ssl_verify: Some(SslVerify::Enabled),
user_agent: Some("litellm/test".into()),
..HttpSettings::default()
}
);
});
}
#[test]
fn python_settings_flow_into_the_configured_layer() {
Python::initialize();
Python::attach(|py| {
let layer = configured(&python_settings(
py,
"
ssl_verify='/etc/ssl/corp.pem',
ssl_certificate='/etc/ssl/client.pem',
ssl_security_level='2',
ssl_ecdh_curve='X25519',
force_ipv4=True,
http2=True,
aiohttp_trust_env=True,
disable_aiohttp_trust_env=True,
disable_aiohttp_transport=True,
user_agent='litellm/9.9.9',
",
))
.unwrap();
assert_eq!(
layer,
HttpSettingsLayer {
ssl_verify: Some(SslVerify::CaBundle("/etc/ssl/corp.pem".into())),
ssl_certificate: Some("/etc/ssl/client.pem".into()),
ssl_security_level: Some("2".into()),
ssl_ecdh_curve: Some("X25519".into()),
force_ipv4: Some(true),
http2: Some(true),
aiohttp_trust_env: Some(true),
disable_aiohttp_trust_env: Some(true),
disable_aiohttp_transport: Some(true),
user_agent: Some("litellm/9.9.9".into()),
..HttpSettingsLayer::default()
}
);
});
}
#[test]
fn user_agent_environment_variable_beats_the_python_default() {
Python::initialize();
Python::attach(|py| {
let settings = HttpSettings::from_layers([
HttpSettingsLayer::from_environment(&|name| {
(name == "LITELLM_USER_AGENT").then(|| "operator/1".to_string())
}),
configured(&python_settings(py, "")).unwrap(),
]);
assert_eq!(settings.user_agent.as_deref(), Some("operator/1"));
});
}
#[rstest]
#[case::disabled("ssl_verify=False", Verify::Disabled)]
#[case::disabled_string("ssl_verify='False'", Verify::Disabled)]
#[case::enabled_string("ssl_verify='true'", Verify::BuiltInRoots)]
#[case::bundle("ssl_verify='/tmp/ca.pem'", Verify::CaBundle("/tmp/ca.pem".into()))]
fn ssl_verify_global_resolves_like_get_ssl_verify(
#[case] overrides: &str,
#[case] expected: Verify,
) {
Python::initialize();
Python::attach(|py| {
let layer = configured(&python_settings(py, overrides)).unwrap();
let config = Resolution::from(&HttpSettings::from_layers([layer])).config;
assert_eq!(config.verify, expected);
});
}
#[test]
fn ssl_context_global_is_ignored_so_environment_and_defaults_apply() {
Python::initialize();
Python::attach(|py| {
let layer = configured(&python_settings(py, "ssl_verify=object()")).unwrap();
assert_eq!(layer.ssl_verify, None);
});
}
#[test]
fn unsupported_settings_are_reported_once_per_process() {
let reported = Mutex::default();
let curve = Unsupported::EcdhCurve("secp521r1".into());
let level = Unsupported::SecurityLevel("@SECLEVEL=1".into());
assert_eq!(
unreported(&reported, vec![curve.clone(), level.clone()]),
[curve.clone(), level]
);
assert_eq!(unreported(&reported, vec![curve]), []);
}
#[test]
fn mistyped_python_settings_decline_instead_of_raising() {
Python::initialize();
Python::attach(|py| {
let error = configured(&python_settings(py, "force_ipv4='yes'")).unwrap_err();
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
});
}
fn configured_ssl_verify(ssl_verify: SslVerify) -> HttpSettingsLayer {
HttpSettingsLayer {
ssl_verify: Some(ssl_verify),
..HttpSettingsLayer::default()
}
}
#[test]
fn call_ssl_verify_beats_the_configured_value() {
Python::initialize();
Python::attach(|py| {
let kwargs = PyDict::new(py);
kwargs.set_item("ssl_verify", false).unwrap();
let call = for_call(call_ssl_verify(&kwargs).unwrap(), true);
let settings =
HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Enabled)]);
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
});
}
#[test]
fn absent_call_ssl_verify_keeps_the_configured_value() {
Python::initialize();
Python::attach(|py| {
let kwargs = PyDict::new(py);
kwargs.set_item("ssl_verify", py.None()).unwrap();
let call = for_call(call_ssl_verify(&kwargs).unwrap(), true);
let settings =
HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]);
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
});
}
#[test]
fn live_ssl_context_argument_is_ignored_so_the_configured_value_applies() {
Python::initialize();
Python::attach(|py| {
let kwargs = PyDict::new(py);
kwargs
.set_item("ssl_verify", py.eval(c"object()", None, None).unwrap())
.unwrap();
let call = for_call(call_ssl_verify(&kwargs).unwrap(), true);
let settings =
HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]);
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
});
}
#[rstest]
#[case::asynchronous(true, false)]
#[case::synchronous(false, true)]
fn synchronous_calls_honor_environment_proxies_even_when_aiohttp_opts_out(
#[case] asynchronous: bool,
#[case] expected: bool,
) {
let opted_out = HttpSettingsLayer {
disable_aiohttp_trust_env: Some(true),
disable_aiohttp_transport: Some(false),
..HttpSettingsLayer::default()
};
let settings = HttpSettings::from_layers([for_call(None, asynchronous), opted_out]);
assert_eq!(settings.trust_proxy_env, expected);
}
}

2
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -1,7 +1,9 @@
mod credentials;
mod diagnostics;
mod errors;
mod http;
mod marshal;
mod python_settings;
mod routes;
mod token_counter;

65
litellm-rust/crates/python-bridge/src/python_settings.rs Normal file
View file

@ -0,0 +1,65 @@
use pyo3::prelude::*;
const MODULE: &str = "litellm.rust_bridge.settings";
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum PythonSettings {
Http,
UrlPolicy,
}
impl PythonSettings {
#[cfg(test)]
pub(crate) const ALL: [Self; 2] = [Self::Http, Self::UrlPolicy];
pub(crate) fn name(self) -> &'static str {
match self {
Self::Http => "http_settings",
Self::UrlPolicy => "url_policy",
}
}
pub(crate) fn read(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
py.import(MODULE)?.getattr(self.name())?.call0()
}
pub(crate) fn warn(py: Python<'_>, message: &str) -> PyResult<()> {
py.import(MODULE)?.getattr("warn")?.call1((message,))?;
Ok(())
}
}
#[cfg(test)]
pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
#[cfg(test)]
mod tests {
use std::{collections::BTreeSet, ffi::CString};
use pyo3::{prelude::*, types::PyDict};
use super::{CONTRACT, PythonSettings};
#[test]
fn every_settings_group_is_in_the_python_contract() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
locals.set_item("contract", CONTRACT).unwrap();
let source = CString::new("import json\nkeys = list(json.loads(contract))").unwrap();
py.run(&source, Some(&locals), Some(&locals)).unwrap();
let declared: BTreeSet<String> = locals
.get_item("keys")
.unwrap()
.unwrap()
.extract::<Vec<String>>()
.unwrap()
.into_iter()
.collect();
let read: BTreeSet<String> = PythonSettings::ALL
.map(|group| group.name().to_owned())
.into();
assert_eq!(read, declared);
});
}
}

16
litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs
View file

@ -3,7 +3,10 @@ mod errors;
mod host;
mod project;
use std::sync::LazyLock;
use host::OcrRouteHost;
use litellm_auth_gcp::VertexAuth;
use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call};
use litellm_core::ocr::route::ocr_machine;
use litellm_llms::custom_httpx::llm_http_handler::OcrClient;
@ -12,6 +15,8 @@ use pyo3::{
types::{PyDict, PyTuple},
};
use crate::{errors::RustBridgeDeclined, http};
const SURFACE: LegacySurface = LegacySurface {
call_type: "ocr",
input_description: "OCR document processing",
@ -23,6 +28,8 @@ const ASYNC_SURFACE: LegacySurface = LegacySurface {
..SURFACE
};
static VERTEX_AUTH: LazyLock<VertexAuth> = LazyLock::new(VertexAuth::default);
fn run_ocr(
py: Python<'_>,
request: Bound<'_, PyAny>,
@ -30,7 +37,14 @@ fn run_ocr(
kwargs: Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let client = OcrClient::shared().map_err(errors::to_pyerr)?;
let config = http::call_config(py, &kwargs, asynchronous)?;
let client = OcrClient::new(
http::pool(),
&config,
http::url_policy(py)?,
VERTEX_AUTH.clone(),
)
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?;
run_legacy_call(
py,
if asynchronous { ASYNC_SURFACE } else { SURFACE },

6
litellm/constants.py
View file

@ -578,6 +578,9 @@ FIREWORKS_AI_176_B_MOE: Final = int(os.getenv("FIREWORKS_AI_176_B_MOE", 176))
FIREWORKS_AI_4_B: Final = int(os.getenv("FIREWORKS_AI_4_B", 4))
FIREWORKS_AI_16_B: Final = int(os.getenv("FIREWORKS_AI_16_B", 16))
FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80))
# https://docs.fireworks.ai/guides/prompt-caching (accessed 2026-09-19): serverless cached prompt tokens
# default to a 50% discount off the input rate
FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO: Final = 0.5
#### Logging callback constants ####
REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM"
MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50))
@ -998,6 +1001,9 @@ openai_compatible_providers: Final[list] = [
"cognition",
"scx-ai",
]
OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS: Final = frozenset({"openai"} | frozenset(openai_compatible_providers))
openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions`
"together_ai",
"fireworks_ai",

9
litellm/exceptions.py
View file

@ -464,6 +464,7 @@ class RateLimitError(openai.RateLimitError):
rate_limit_type: str | RateLimitType | None = None,
headers: dict[str, str] | None = None,
detail: Any = None,
body: object | None = None,
):
self.status_code = 429
self.message = f"litellm.RateLimitError: {message}"
@ -507,7 +508,7 @@ class RateLimitError(openai.RateLimitError):
),
)
super().__init__(
self.message, response=self.response, body=None
self.message, response=self.response, body=body
) # Call the base class constructor with the parameters it needs
self.code = "429"
self.type = "throttling_error"
@ -765,6 +766,7 @@ class InternalServerError(openai.InternalServerError):
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
body: object | None = None,
):
self.status_code = 500
self.message = f"litellm.InternalServerError: {message}"
@ -783,7 +785,7 @@ class InternalServerError(openai.InternalServerError):
),
)
super().__init__(
self.message, response=self.response, body=None
self.message, response=self.response, body=body
) # Call the base class constructor with the parameters it needs
def __str__(self):
@ -815,6 +817,7 @@ class APIError(openai.APIError):
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
body: object | None = None,
):
self.status_code = status_code
self.message = f"litellm.APIError: {message}"
@ -825,7 +828,7 @@ class APIError(openai.APIError):
self.num_retries = num_retries
if request is None:
request = httpx.Request(method="POST", url="https://api.openai.com/v1")
super().__init__(self.message, request=request, body=None)
super().__init__(self.message, request=request, body=body)
def __str__(self):
_message = self.message

2
litellm/experimental_mcp_client/Readme.md
View file

@ -6,7 +6,7 @@ LiteLLM MCP Client allows you to use MCP tools with LiteLLM
The `mcp` and `proxy` extras require MCP Python SDK 2.2 or newer within the 2.x release line. Installing core LiteLLM without these extras does not require MCP
Existing MCP SDK1 clients can continue connecting to the gateway over the supported legacy MCP protocols. The client and gateway can use different SDK versions in separate Python environments. Modern protocol advertisement remains disabled during the Phase 0 upgrade
Existing MCP SDK1 clients can continue connecting to the gateway over the supported legacy MCP protocols. The client and gateway can use different SDK versions in separate Python environments. Modern protocol advertisement remains disabled during the Phase 0 upgrade. An initialize body requesting `2026-07-28` falls back to the supported legacy version `2025-11-25`; an explicit `MCP-Protocol-Version: 2026-07-28` HTTP header is rejected with HTTP 400
Code sharing the gateway's Python environment must support SDK2. Its Python API has breaking changes, including renamed imports and snake_case model attributes such as `input_schema`, `is_error`, and `structured_content`. This also applies to callers consuming SDK objects returned by LiteLLM's experimental MCP client. MCP JSON fields retain their protocol spelling, such as `inputSchema` and `isError`

3
litellm/experimental_mcp_client/client.py
View file

@ -159,8 +159,7 @@ class _MCPHTTPClient(httpx2.AsyncClient):
follow_redirects: bool | UseClientDefault = httpx2.USE_CLIENT_DEFAULT,
) -> httpx2.Response:
response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects)
# Check after the auth flow completes so a refreshable 401 can still be retried.
if request.method == "POST" and response.is_error:
if request.method == "POST" and response.is_error and response.status_code != 404:
await response.aclose()
response.raise_for_status()
return response

106
litellm/integrations/websearch_interception/handler.py
View file

@ -44,6 +44,8 @@ from litellm.types.integrations.custom_logger import (
from litellm.types.integrations.websearch_interception import (
AnthropicSearchQuery,
AnthropicServerToolUseBlock,
SearchFailed,
SearchOutcome,
WebSearchInterceptionConfig,
)
from litellm.types.llms.anthropic import AnthropicThinkingParam
@ -332,16 +334,8 @@ class WebSearchInterceptionLogger(CustomLogger):
None,
)
# Execute search — keep the structured SearchResponse so the native
# block can carry per-result url/title/page_age.
try:
if kwargs is None:
search_result_text, structured = await self._execute_search(query)
else:
search_result_text, structured = await self._execute_search(query, kwargs=kwargs)
except Exception as e:
verbose_logger.error("WebSearchInterception: Short-circuit search failed: %s", e)
search_result_text, structured = f"Search failed: {e}", None
outcome: Final = await self._short_circuit_search_outcome(query, kwargs=kwargs)
search_result_text: Final = WebSearchTransformation.search_outcome_text(outcome)
content: Final[list[dict[str, object]]] = []
if native_tool is not None:
@ -356,10 +350,7 @@ class WebSearchInterceptionLogger(CustomLogger):
}
)
content.append(
WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id=tool_use_id,
search_response=structured,
)
WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome)
)
# Keep the text block so non-native short-circuit callers (Claude Code,
# github_copilot, etc.) see the same payload they always have.
@ -934,7 +925,7 @@ class WebSearchInterceptionLogger(CustomLogger):
tool_calls: Final = tools["tool_calls"]
thinking_blocks: Final = tools.get("thinking_blocks", [])
request_patch, structured_results = await self._build_anthropic_request_patch(
request_patch, search_outcomes = await self._build_anthropic_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
@ -953,17 +944,21 @@ class WebSearchInterceptionLogger(CustomLogger):
# pre-build the Anthropic-native ``web_search_tool_result`` blocks now
# (while we still have the structured SearchResponse list) and stash
# them on plan metadata for the post-hook to inject.
if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY):
metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks(
tool_calls=tool_calls,
structured_results=structured_results,
)
if not kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY):
return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata)
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=request_patch,
metadata=metadata,
metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks(
tool_calls=tool_calls,
search_outcomes=search_outcomes,
)
every_search_failed: Final = bool(search_outcomes) and all(
isinstance(outcome, SearchFailed) for outcome in search_outcomes
)
if every_search_failed:
return AgenticLoopPlan(
run_agentic_loop=False, terminate=True, stop_reason="web_search_failed", metadata=metadata
)
return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata)
async def async_post_agentic_loop_response_hook(
self,
@ -992,7 +987,7 @@ class WebSearchInterceptionLogger(CustomLogger):
@staticmethod
def _build_native_result_blocks(
tool_calls: list[dict],
structured_results: list[SearchResponse | None],
search_outcomes: Sequence[SearchOutcome],
) -> tuple[Mapping[str, object], ...]:
"""
Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call.
@ -1004,10 +999,10 @@ class WebSearchInterceptionLogger(CustomLogger):
"""
return tuple(
block
for i, tool_call in enumerate(tool_calls)
for tool_call, outcome in zip(tool_calls, search_outcomes, strict=True)
for block in WebSearchInterceptionLogger._native_result_pair(
query=WebSearchInterceptionLogger._tool_call_query(tool_call),
search_response=structured_results[i] if i < len(structured_results) else None,
outcome=outcome,
)
)
@ -1022,15 +1017,12 @@ class WebSearchInterceptionLogger(CustomLogger):
@staticmethod
def _native_result_pair(
query: str,
search_response: SearchResponse | None,
outcome: SearchOutcome,
) -> tuple[Mapping[str, object], Mapping[str, object]]:
tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}"
return (
AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(),
WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id=tool_use_id,
search_response=search_response,
),
WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome),
)
@staticmethod
@ -1306,7 +1298,7 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs: Mapping[str, object],
) -> "AnthropicMessagesResponse | AsyncIterator[object]":
"""Legacy path: execute search + build patch + run follow-up call."""
request_patch, structured_results = await self._build_anthropic_request_patch(
request_patch, search_outcomes = await self._build_anthropic_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
@ -1344,7 +1336,7 @@ class WebSearchInterceptionLogger(CustomLogger):
if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY):
native_blocks: Final = self._build_native_result_blocks(
tool_calls=tool_calls,
structured_results=structured_results,
search_outcomes=search_outcomes,
)
response = self._inject_native_blocks(response, native_blocks)
@ -1359,15 +1351,9 @@ class WebSearchInterceptionLogger(CustomLogger):
anthropic_messages_optional_request_params: dict,
logging_obj: "LiteLLMLoggingObj | None",
kwargs: dict,
) -> tuple[AgenticLoopRequestPatch, list[SearchResponse | None]]:
) -> tuple[AgenticLoopRequestPatch, tuple[SearchOutcome, ...]]:
"""
Execute litellm.search() and build follow-up request patch.
Returns the patch alongside the parallel list of structured
``SearchResponse`` objects (one per tool_call, ``None`` when the
search failed or the tool_call had no query). The caller uses these
to optionally build Anthropic-native ``web_search_tool_result``
content blocks for the final response.
"""
# Extract search queries from tool_use blocks
@ -1385,27 +1371,10 @@ class WebSearchInterceptionLogger(CustomLogger):
# Execute searches in parallel
verbose_logger.debug("WebSearchInterception: Executing %s search(es) in parallel", len(search_tasks))
search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True)
# Split the gathered (text, structured) tuples into two parallel lists.
# The text list feeds the follow-up model call; the structured list
# is returned to the caller for native-block emission.
final_search_results: Final[list[str]] = []
structured_results: Final[list[SearchResponse | None]] = []
for i, result in enumerate(search_results):
if isinstance(result, Exception):
verbose_logger.error("WebSearchInterception: Search %s failed with error: %s", i, result)
final_search_results.append(f"Search failed: {result}")
structured_results.append(None)
elif isinstance(result, tuple) and len(result) == 2:
text_value, structured_value = result
final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value))
structured_results.append(structured_value if isinstance(structured_value, SearchResponse) else None)
else:
# Defensive: legacy callers / unexpected shape — preserve text,
# drop structure.
verbose_logger.debug("WebSearchInterception: Unexpected result type %s at index %s", type(result), i)
final_search_results.append(str(result))
structured_results.append(None)
search_outcomes: Final = tuple(WebSearchTransformation.search_outcome(result) for result in search_results)
final_search_results: Final = tuple(
WebSearchTransformation.search_outcome_text(outcome) for outcome in search_outcomes
)
# Build assistant and user messages using transformation
assistant_message, user_message = WebSearchTransformation.transform_response(
@ -1449,7 +1418,18 @@ class WebSearchInterceptionLogger(CustomLogger):
optional_params=optional_params_without_max_tokens,
kwargs=kwargs_for_followup,
)
return patch, structured_results
return patch, search_outcomes
async def _short_circuit_search_outcome(self, query: str, kwargs: Mapping[str, object] | None) -> SearchOutcome:
try:
result: Final = (
await self._execute_search(query)
if kwargs is None
else await self._execute_search(query, kwargs=kwargs)
)
except Exception as e:
return WebSearchTransformation.search_outcome(e)
return WebSearchTransformation.search_outcome(result)
async def _execute_search(
self, query: str, kwargs: Mapping[str, object] | None = None

77
litellm/integrations/websearch_interception/transformation.py
View file

@ -5,11 +5,21 @@ Transforms between Anthropic/OpenAI tool_use format and LiteLLM search format.
"""
import json
from collections.abc import Sequence
from typing import Any, Final
from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME
from litellm.exceptions import BadRequestError, RateLimitError
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.types.integrations.websearch_interception import (
SearchFailed,
SearchOutcome,
SearchSucceeded,
WebSearchToolResultErrorCode,
)
class WebSearchTransformation:
@ -280,7 +290,7 @@ class WebSearchTransformation:
@staticmethod
def transform_response(
tool_calls: list[dict],
search_results: list[str],
search_results: Sequence[str],
response_format: str = "anthropic",
thinking_blocks: list[dict] | None = None,
) -> tuple[dict, dict | list[dict]]:
@ -314,7 +324,7 @@ class WebSearchTransformation:
@staticmethod
def _transform_response_anthropic(
tool_calls: list[dict],
search_results: list[str],
search_results: Sequence[str],
thinking_blocks: list[dict] | None = None,
) -> tuple[dict, dict]:
"""Transform to Anthropic format (single user message with tool_result blocks)"""
@ -364,7 +374,7 @@ class WebSearchTransformation:
@staticmethod
def _transform_response_openai(
tool_calls: list[dict],
search_results: list[str],
search_results: Sequence[str],
) -> tuple[dict, list[dict]]:
"""Transform to OpenAI format (assistant with tool_calls, separate tool messages)"""
# Build assistant message with tool_calls
@ -456,6 +466,67 @@ class WebSearchTransformation:
"content": items,
}
@staticmethod
def build_web_search_tool_result_error_block(
tool_use_id: str,
error_code: WebSearchToolResultErrorCode,
) -> dict[str, object]:
return {
"type": "web_search_tool_result",
"tool_use_id": tool_use_id,
"content": {"type": "web_search_tool_result_error", "error_code": error_code},
}
@staticmethod
def build_web_search_outcome_block(tool_use_id: str, outcome: SearchOutcome) -> dict[str, object]:
match outcome:
case SearchSucceeded(response=response):
return WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id=tool_use_id,
search_response=response,
)
case SearchFailed(error_code=error_code):
return WebSearchTransformation.build_web_search_tool_result_error_block(
tool_use_id=tool_use_id,
error_code=error_code,
)
case _:
assert_never(outcome)
@staticmethod
def search_error_code(error: BaseException) -> WebSearchToolResultErrorCode:
match error:
case RateLimitError():
return "too_many_requests"
case BadRequestError():
return "invalid_tool_input"
case _:
return "unavailable"
@staticmethod
def search_outcome(result: object) -> SearchOutcome:
match result:
case BaseException():
verbose_logger.error("WebSearchInterception: Search failed with error: %s", result)
return SearchFailed(error_code=WebSearchTransformation.search_error_code(result), message=str(result))
case (str() as text, SearchResponse() as response):
return SearchSucceeded(text=text, response=response)
case (str() as text, None):
return SearchSucceeded(text=text, response=None)
case _:
verbose_logger.debug("WebSearchInterception: Unexpected search result type %s", type(result))
return SearchSucceeded(text=str(result), response=None)
@staticmethod
def search_outcome_text(outcome: SearchOutcome) -> str:
match outcome:
case SearchSucceeded(text=text):
return text
case SearchFailed(message=message):
return f"Search failed: {message}"
case _:
assert_never(outcome)
@staticmethod
def format_search_response(result: SearchResponse) -> str:
"""

5
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -307,6 +307,7 @@ def _map_openai_exception(
model=model,
llm_provider=custom_llm_provider,
response=response,
body=getattr(original_exception, "body", None),
)
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
raise ContextWindowExceededError(
@ -381,6 +382,7 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
body=getattr(original_exception, "body", None),
)
elif "Request too large" in error_str:
raise RateLimitError(
@ -389,6 +391,7 @@ def _map_openai_exception(
llm_provider=custom_llm_provider,
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
elif (
"The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable"
@ -460,6 +463,7 @@ def _map_openai_exception(
llm_provider=custom_llm_provider,
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
elif original_exception.status_code == 500:
raise InternalServerError(
@ -468,6 +472,7 @@ def _map_openai_exception(
llm_provider=custom_llm_provider,
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
elif original_exception.status_code == 502:
raise BadGatewayError(

37
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -18,6 +18,7 @@ from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import (
select_tier_for_input,
tier_rate,
)
from litellm.llms.fireworks_ai.cache_pricing import with_default_cache_read_rate
from litellm.types.utils import (
CacheCreationTokenDetails,
CallTypes,
@ -72,6 +73,12 @@ def _uses_inclusive_token_thresholds(custom_llm_provider: str | None) -> bool:
return custom_llm_provider in _INCLUSIVE_THRESHOLD_PROVIDERS
def apply_provider_cache_read_default(model_info: ModelInfo, custom_llm_provider: str | None) -> ModelInfo:
if custom_llm_provider == "fireworks_ai":
return with_default_cache_read_rate(model_info)
return model_info
def _get_token_detail_value(details: object, key: str) -> int | None:
if isinstance(details, dict):
value = details.get(key)
@ -1170,8 +1177,10 @@ def generic_cost_per_token(
# rather than handing back a name for this to re-resolve. A name cannot express a
# per-deployment override: those are registered under the deployment id and kept off
# the shared model-name key, so resolving from the name here reads the public rate.
if model_info is None:
model_info = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
resolved_model_info: Final = apply_provider_cache_read_default(
get_model_info(model=model, custom_llm_provider=custom_llm_provider) if model_info is None else model_info,
custom_llm_provider,
)
## CALCULATE INPUT COST
### Cost of processing (non-cache hit + cache hit) + Cost of cache-writing (cache writing)
@ -1236,7 +1245,7 @@ def generic_cost_per_token(
cache_creation_cost_above_1hr,
cache_read_cost,
) = _get_token_base_cost(
model_info=model_info,
model_info=resolved_model_info,
usage=usage,
service_tier=service_tier,
current_time=billing_time,
@ -1245,7 +1254,7 @@ def generic_cost_per_token(
prompt_cost = _calculate_input_cost(
prompt_tokens_details=prompt_tokens_details,
model_info=model_info,
model_info=resolved_model_info,
prompt_base_cost=prompt_base_cost,
cache_read_cost=cache_read_cost,
cache_creation_cost=cache_creation_cost,
@ -1290,7 +1299,7 @@ def generic_cost_per_token(
## AUDIO COST
if not is_text_tokens_total and audio_tokens is not None and audio_tokens > 0:
_output_cost_per_audio_token = _get_cost_per_unit(model_info, "output_cost_per_audio_token", None)
_output_cost_per_audio_token = _get_cost_per_unit(resolved_model_info, "output_cost_per_audio_token", None)
_output_cost_per_audio_token = (
_output_cost_per_audio_token if _output_cost_per_audio_token is not None else completion_base_cost
)
@ -1299,7 +1308,7 @@ def generic_cost_per_token(
## REASONING COST
if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0:
completion_cost += float(reasoning_tokens) * _resolve_billed_reasoning_rate(
model_info=model_info,
model_info=resolved_model_info,
usage=usage,
service_tier=service_tier,
completion_base_cost=completion_base_cost,
@ -1308,7 +1317,7 @@ def generic_cost_per_token(
## IMAGE COST
if not is_text_tokens_total and image_tokens and image_tokens > 0:
_output_cost_per_image_token = _get_cost_per_unit(model_info, "output_cost_per_image_token", None)
_output_cost_per_image_token = _get_cost_per_unit(resolved_model_info, "output_cost_per_image_token", None)
_output_cost_per_image_token = (
_output_cost_per_image_token if _output_cost_per_image_token is not None else completion_base_cost
)
@ -1316,7 +1325,7 @@ def generic_cost_per_token(
## VIDEO COST
if not is_text_tokens_total and video_tokens and video_tokens > 0:
_output_cost_per_video_token = _get_cost_per_unit(model_info, "output_cost_per_video_token", None)
_output_cost_per_video_token = _get_cost_per_unit(resolved_model_info, "output_cost_per_video_token", None)
_output_cost_per_video_token = (
_output_cost_per_video_token if _output_cost_per_video_token is not None else completion_base_cost
)
@ -1325,12 +1334,12 @@ def generic_cost_per_token(
## REGIONAL DATA-RESIDENCY UPLIFT
# Applied as a flat multiplier across all token costs for the request
# when the upstream is a regionalized OpenAI host (eu./us.api.openai.com).
uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency)
uplift: Final = _get_regional_uplift_multiplier(resolved_model_info, data_residency)
if uplift != 1.0:
prompt_cost *= uplift
completion_cost *= uplift
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(resolved_model_info, vertex_location)
if vertex_uplift != 1.0:
prompt_cost *= vertex_uplift
completion_cost *= vertex_uplift
@ -1487,7 +1496,10 @@ def get_billed_token_rates(
if custom_cost_per_token is not None:
return _custom_pricing_rates(custom_cost_per_token)
try:
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
model_info: Final = apply_provider_cache_read_default(
get_model_info(model=model, custom_llm_provider=custom_llm_provider),
custom_llm_provider,
)
except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates
return None
return _cost_map_billed_rates(
@ -1578,8 +1590,9 @@ def calculate_prompt_caching_savings(
``billed_at`` is the request's completion time, so off-peak windows resolve as the
biller saw them rather than at the later spend write.
"""
model_info_with_cache_read_default: Final = apply_provider_cache_read_default(model_info, custom_llm_provider)
prompt_base_cost, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = _get_token_base_cost(
model_info=model_info,
model_info=model_info_with_cache_read_default,
usage=usage,
service_tier=service_tier,
current_time=billed_at,

63
litellm/litellm_core_utils/llm_response_utils/response_metadata.py
View file

@ -1,11 +1,12 @@
import datetime
from collections.abc import Mapping
from functools import reduce
from typing import Any, Final
import httpx
from litellm.constants import LITELLM_DETAILED_TIMING
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
from litellm.litellm_core_utils.logging_utils import LiteLLMLoggingObject
from litellm.types.utils import (
@ -16,19 +17,59 @@ from litellm.types.utils import (
)
def _timing_window_start(
start_time: datetime.datetime, logging_obj: LiteLLMLoggingObject
) -> tuple[datetime.datetime, bool]:
received_at: Final = get_litellm_metadata_from_kwargs(logging_obj.model_call_details).get("litellm_received_at")
if isinstance(received_at, datetime.datetime):
return received_at, True
return start_time, False
def _union_duration_ms(windows: object, lower: float, upper: float) -> float | None:
if not isinstance(windows, (list, tuple)):
return None
clipped: Final[tuple[tuple[float, float], ...]] = tuple(
(max(lower, float(window[0])), min(upper, float(window[1])))
for window in windows
if isinstance(window, (list, tuple))
and len(window) == 2
and isinstance(window[0], (int, float))
and isinstance(window[1], (int, float))
and max(lower, float(window[0])) < min(upper, float(window[1]))
)
if not clipped:
return None
ordered: Final[tuple[tuple[float, float], ...]] = tuple(sorted(clipped))
def merge_window(
merged: tuple[tuple[float, float], ...], current: tuple[float, float]
) -> tuple[tuple[float, float], ...]:
if not merged or current[0] > merged[-1][1]:
return (*merged, current)
return (*merged[:-1], (merged[-1][0], max(merged[-1][1], current[1])))
merged: Final[tuple[tuple[float, float], ...]] = reduce(merge_window, ordered, ())
return sum(end - start for start, end in merged) * 1000
def response_timing_metrics(
start_time: datetime.datetime,
end_time: datetime.datetime,
logging_obj: LiteLLMLoggingObject,
include_overhead: bool = True,
) -> Mapping[str, float]:
"""``_response_ms`` for the whole call, plus ``litellm_overhead_time_ms`` when it can be derived.
"""``_response_ms`` for the window starting at proxy receive time when stamped, else ``start_time``.
On a cache hit the overhead is the total minus the cache read; otherwise it is the total minus
the provider call (``llm_api_duration_ms``). It is omitted when neither duration was recorded,
and when ``include_overhead`` is False because the two durations cover different windows.
"""
total_response_time_ms: Final = (end_time - start_time).total_seconds() * 1000
timing_window: Final = _timing_window_start(start_time, logging_obj)
window_start: Final = timing_window[0]
receive_anchored: Final = timing_window[1]
total_response_time_ms: Final = (end_time.timestamp() - window_start.timestamp()) * 1000
if not include_overhead:
return {"_response_ms": total_response_time_ms} # mutable-ok: read-only timing result
caching_details: Final = logging_obj.caching_details
@ -37,11 +78,22 @@ def response_timing_metrics(
if caching_details is not None and caching_details.get("cache_hit") is True
else None
)
metadata: Final[Mapping[str, object]] = get_litellm_metadata_from_kwargs(logging_obj.model_call_details)
llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms")
if cache_duration_ms is not None:
overhead_ms: float | None = total_response_time_ms - cache_duration_ms
elif llm_api_duration_ms is not None:
overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4)
provider_duration_ms: Final[float | None] = (
_union_duration_ms(
metadata.get("llm_api_timing_windows"),
window_start.timestamp(),
end_time.timestamp(),
)
if receive_anchored
else None
)
effective: Final = provider_duration_ms if provider_duration_ms is not None else llm_api_duration_ms
overhead_ms = round(total_response_time_ms - effective, 4) if isinstance(effective, (int, float)) else None
else:
overhead_ms = None
if overhead_ms is None:
@ -152,7 +204,8 @@ class ResponseMetadata:
# pre-processing = time from request start to LLM API call start
api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time")
if api_call_start is not None and start_time is not None:
pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000
anchor: Final = _timing_window_start(start_time, logging_obj)[0]
pre_ms: Final = (api_call_start.timestamp() - anchor.timestamp()) * 1000
detailed["timing_pre_processing_ms"] = round(pre_ms, 4)
# post-processing = total - pre - llm_api

15
litellm/litellm_core_utils/logging_utils.py
View file

@ -12,6 +12,7 @@ from litellm.constants import (
BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS,
MAX_BASE64_LENGTH_FOR_LOGGING,
)
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -286,6 +287,20 @@ def _set_duration_in_model_call_details(
duration_ms: Final = (end_time - start_time).total_seconds() * 1000
if logging_obj and hasattr(logging_obj, "model_call_details"):
logging_obj.model_call_details["llm_api_duration_ms"] = duration_ms
metadata: Final[dict[str, object]] = get_litellm_metadata_from_kwargs(logging_obj.model_call_details)
recorded: Final = metadata.get("llm_api_timing_windows")
earlier: Final[tuple[tuple[float, float], ...]] = tuple(
(float(window[0]), float(window[1]))
for window in (recorded if isinstance(recorded, (list, tuple)) else ())
if isinstance(window, (list, tuple))
and len(window) == 2
and isinstance(window[0], (int, float))
and isinstance(window[1], (int, float))
)
metadata["llm_api_timing_windows"] = (
*earlier,
(start_time.timestamp(), end_time.timestamp()),
)
else:
verbose_logger.debug("`logging_obj` not found - unable to track `llm_api_duration_ms")
except Exception as e:

4
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -12,7 +12,7 @@ import litellm
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend
from litellm.types.llms.openai import (
OpenAIRealtimeEvents,
OpenAIRealtimeOutputItemDone,
@ -127,7 +127,7 @@ class RealTimeStreaming:
def __init__(
self,
websocket: Any,
backend_ws: CLIENT_CONNECTION_CLASS,
backend_ws: CLIENT_CONNECTION_CLASS | RealtimeBackend,
logging_obj: LiteLLMLogging,
provider_config: BaseRealtimeConfig | None = None,
model: str = "",

22
litellm/llms/anthropic/common_utils.py
View file

@ -1435,12 +1435,19 @@ class _ReplayedWebSearchResult(BaseModel):
encrypted_content: str = ""
class _ReplayedWebSearchToolResultError(BaseModel):
model_config = ConfigDict(extra="allow")
type: Literal["web_search_tool_result_error"]
error_code: str = ""
class _ReplayedWebSearchToolResult(BaseModel):
model_config = ConfigDict(extra="allow")
type: Literal["web_search_tool_result"]
tool_use_id: str
content: tuple[_ReplayedWebSearchResult, ...]
content: tuple[_ReplayedWebSearchResult, ...] | _ReplayedWebSearchToolResultError
class _ReplayedServerToolUse(BaseModel):
@ -1464,17 +1471,12 @@ def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchTool
"""
The parsed block when it is a ``web_search_tool_result`` carrying no
``encrypted_content``, else None for anything Anthropic itself issued.
An empty ``content`` list is flattenable too. It is what the interceptor emits
when a search legitimately returns nothing and when a search raises, and it
carries neither evidence to preserve nor an ``encrypted_content`` to respect,
so leaving it in place only buys the 400 this whole function exists to avoid.
"""
try:
parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block)
except ValidationError:
return None
if any(result.encrypted_content for result in parsed.content):
if isinstance(parsed.content, tuple) and any(result.encrypted_content for result in parsed.content):
return None
return parsed
@ -1486,8 +1488,12 @@ def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None:
return None
def _render_web_search_results(query: str, results: tuple[_ReplayedWebSearchResult, ...]) -> str:
def _render_web_search_results(
query: str, results: tuple[_ReplayedWebSearchResult, ...] | _ReplayedWebSearchToolResultError
) -> str:
header: Final = f"Web search results for '{query}':" if query else "Web search results:"
if isinstance(results, _ReplayedWebSearchToolResultError):
return f"{header}\n\nSearch failed: {results.error_code or 'unavailable'}"
if not results:
return f"{header}\n\nNo results were returned."
body: Final = "\n\n".join(

9
litellm/llms/base_llm/audio_transcription/transformation.py
View file

@ -52,6 +52,15 @@ class BaseAudioTranscriptionConfig(BaseConfig, ABC):
"""
return False
@property
def has_native_transcription_endpoint(self) -> bool:
"""
Opt-in for OpenAI-compatible providers whose transcription lives on a
non-OpenAI route: when True the request skips the OpenAI SDK transport
and goes through this config via the shared http handler.
"""
return False
def get_complete_url(
self,
api_base: str | None,

275
litellm/llms/base_llm/realtime/transcription_protocol.py Normal file
View file

@ -0,0 +1,275 @@
import base64
import binascii
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm._uuid import uuid
from litellm.types.llms.openai import (
OpenAIRealtimeErrorEvent,
OpenAIRealtimeInputAudioBufferSpeechEvent,
OpenAIRealtimeInputAudioTranscriptionCompleted,
OpenAIRealtimeInputAudioTranscriptionDelta,
OpenAIRealtimeServerVadTurnDetection,
OpenAIRealtimeTranscriptionSession,
OpenAIRealtimeTranscriptionSessionCreated,
OpenAIRealtimeTranscriptionSettings,
)
from litellm.types.realtime import RealtimeInputAudioTranscriptionDurationUsage, RealtimeInputAudioTranscriptionUsage
SESSION_UPDATE_EVENT_TYPES: Final = frozenset(("session.update", "transcription_session.update"))
PCM16_ENCODINGS: Final = frozenset(("pcm16", "audio/pcm"))
SERVER_VAD_TURN_DETECTION: Final[OpenAIRealtimeServerVadTurnDetection] = {"type": "server_vad"}
EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({})
_SUPPORTED_TRANSCRIPTION_KEYS: Final = frozenset(("model", "language"))
_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
class RealtimeTranscriptionProtocolError(ValueError):
pass
@dataclass(frozen=True, slots=True)
class TranscriptionAudioFormat:
layout: Literal["beta", "ga"]
encoding: str | None
rate: int | None
channels: int | None
@property
def is_pcm16(self) -> bool:
return self.encoding in PCM16_ENCODINGS
@dataclass(frozen=True, slots=True)
class TranscriptionSessionUpdate:
session_type: str | None
audio_format: TranscriptionAudioFormat | None
model: str | None
language: str | None
unsupported_transcription_keys: tuple[str, ...]
turn_detection: Mapping[str, JsonValue] | None
turn_detection_disabled: bool
@property
def turn_detection_type(self) -> JsonValue | None:
return None if self.turn_detection is None else self.turn_detection.get("type")
ProtocolErrorType = type[RealtimeTranscriptionProtocolError]
def json_object(payload: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError) -> Mapping[str, JsonValue]:
try:
value: Final = _JSON_ADAPTER.validate_json(payload)
except ValidationError:
raise error("invalid JSON object") from None
if not isinstance(value, dict):
raise error("message must be a JSON object")
return value
def json_mapping(
value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError
) -> Mapping[str, JsonValue]:
if value is None:
return EMPTY_JSON_OBJECT
if not isinstance(value, dict):
raise error(f"{name} must be an object")
return value
def json_string(
value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError
) -> str | None:
if value is None:
return None
if not isinstance(value, str):
raise error(f"{name} must be a string")
return value
def json_integer(
value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError
) -> int | None:
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int):
raise error(f"{name} must be an integer")
return value
def new_event_id() -> str:
return f"event_{uuid.uuid4().hex}"
def parse_transcription_session_update(
payload: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError
) -> TranscriptionSessionUpdate:
message: Final = json_object(payload, error)
if message.get("type") not in SESSION_UPDATE_EVENT_TYPES:
raise error("expected session.update")
session: Final = json_mapping(message.get("session"), "session", error)
if not session:
raise error("session.update requires a session object")
audio: Final = json_mapping(session.get("audio"), "session.audio", error)
audio_input: Final = json_mapping(audio.get("input"), "session.audio.input", error)
beta_transcription: Final = session.get("input_audio_transcription")
ga_transcription: Final = audio_input.get("transcription")
if beta_transcription is not None and ga_transcription is not None:
raise error("input transcription must use either beta or GA layout")
transcription: Final = json_mapping(
beta_transcription if beta_transcription is not None else ga_transcription,
"input audio transcription",
error,
)
turn_detection_present: Final = "turn_detection" in session or "turn_detection" in audio_input
turn_detection: Final = session.get("turn_detection", audio_input.get("turn_detection"))
return TranscriptionSessionUpdate(
session_type=json_string(session.get("type"), "session.type", error),
audio_format=_parse_audio_format(session, audio_input, error),
model=json_string(transcription.get("model"), "transcription model", error),
language=json_string(transcription.get("language"), "language", error),
unsupported_transcription_keys=tuple(
sorted(key for key in transcription if key not in _SUPPORTED_TRANSCRIPTION_KEYS)
),
turn_detection=None if turn_detection is None else json_mapping(turn_detection, "turn_detection", error),
turn_detection_disabled=turn_detection_present and turn_detection is None,
)
def _parse_audio_format(
session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue], error: ProtocolErrorType
) -> TranscriptionAudioFormat | None:
beta_format: Final = session.get("input_audio_format")
ga_format: Final = audio_input.get("format")
if beta_format is not None and ga_format is not None:
raise error("input audio format must use either beta or GA layout")
if beta_format is not None:
return TranscriptionAudioFormat(
layout="beta",
encoding=json_string(beta_format, "session.input_audio_format", error),
rate=None,
channels=None,
)
if ga_format is None:
return None
if isinstance(ga_format, str):
return TranscriptionAudioFormat(layout="ga", encoding=ga_format, rate=None, channels=None)
format_mapping: Final = json_mapping(ga_format, "session.audio.input.format", error)
return TranscriptionAudioFormat(
layout="ga",
encoding=json_string(format_mapping.get("type"), "session.audio.input.format.type", error),
rate=json_integer(format_mapping.get("rate"), "session.audio.input.format.rate", error),
channels=json_integer(format_mapping.get("channels"), "session.audio.input.format.channels", error),
)
def decode_pcm16_append(
audio: JsonValue | None,
max_encoded_bytes: int | None = None,
error: ProtocolErrorType = RealtimeTranscriptionProtocolError,
) -> bytes:
if not isinstance(audio, str):
raise error("Audio must be a base64 string")
if max_encoded_bytes is not None and len(audio) > max_encoded_bytes:
raise error("Audio append exceeds the four-second backlog limit")
try:
decoded: Final = base64.b64decode(audio, validate=True)
except (binascii.Error, ValueError):
raise error("Audio must be valid base64") from None
if len(decoded) % 2:
raise error("PCM16 audio must contain complete samples")
return decoded
def _transcription_settings(model: str, language: str | None) -> OpenAIRealtimeTranscriptionSettings:
if language is None:
model_only: Final[OpenAIRealtimeTranscriptionSettings] = {"model": model}
return model_only
with_language: Final[OpenAIRealtimeTranscriptionSettings] = {"model": model, "language": language}
return with_language
def transcription_session(
*, session_id: str, model: str, sample_rate: int, language: str | None, server_vad: bool
) -> OpenAIRealtimeTranscriptionSession:
settings: Final = _transcription_settings(model, language)
session: Final[OpenAIRealtimeTranscriptionSession] = {
"id": session_id,
"object": "realtime.transcription_session",
"type": "transcription",
"audio": {
"input": {
"format": {"type": "audio/pcm", "rate": sample_rate},
"transcription": settings,
"turn_detection": SERVER_VAD_TURN_DETECTION if server_vad else None,
}
},
}
return session
def transcription_session_created_event(
session: OpenAIRealtimeTranscriptionSession,
) -> OpenAIRealtimeTranscriptionSessionCreated:
event: Final[OpenAIRealtimeTranscriptionSessionCreated] = {
"type": "session.created",
"event_id": new_event_id(),
"session": session,
}
return event
def error_event(message: str) -> OpenAIRealtimeErrorEvent:
event: Final[OpenAIRealtimeErrorEvent] = {
"type": "error",
"error": {"type": "server_error", "message": message},
}
return event
def speech_event(
event_type: Literal["input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"], item_id: str
) -> OpenAIRealtimeInputAudioBufferSpeechEvent:
event: Final[OpenAIRealtimeInputAudioBufferSpeechEvent] = {
"type": event_type,
"event_id": new_event_id(),
"item_id": item_id,
}
return event
def delta_event(item_id: str, delta: str) -> OpenAIRealtimeInputAudioTranscriptionDelta:
event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = {
"type": "conversation.item.input_audio_transcription.delta",
"event_id": new_event_id(),
"item_id": item_id,
"content_index": 0,
"delta": delta,
}
return event
def completed_event(
item_id: str, transcript: str, usage: RealtimeInputAudioTranscriptionUsage | None
) -> OpenAIRealtimeInputAudioTranscriptionCompleted:
event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = {
"type": "conversation.item.input_audio_transcription.completed",
"event_id": new_event_id(),
"item_id": item_id,
"content_index": 0,
"transcript": transcript,
}
if usage is None:
return event
billed: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = {**event, "usage": usage}
return billed
def duration_usage(seconds: float) -> RealtimeInputAudioTranscriptionUsage:
usage: Final[RealtimeInputAudioTranscriptionDurationUsage] = {"type": "duration", "seconds": seconds}
return usage

24
litellm/llms/base_llm/realtime/transformation.py
View file

@ -1,8 +1,10 @@
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any
from types import TracebackType
from typing import TYPE_CHECKING, Any, Protocol
import httpx
from typing_extensions import Self
from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
from litellm.types.realtime import (
@ -21,6 +23,23 @@ else:
LiteLLMLoggingObj = Any
class RealtimeBackend(Protocol):
async def __aenter__(self) -> Self: ...
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None: ...
async def send(self, message: str | bytes) -> None: ...
async def recv(self, decode: bool | None = None) -> str | bytes: ...
async def close(self) -> None: ...
class BaseRealtimeConfig(ABC):
@abstractmethod
def validate_environment(
@ -78,6 +97,9 @@ class BaseRealtimeConfig(ABC):
def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None:
return None
async def open_backend(self, url: str, headers: Mapping[str, str]) -> RealtimeBackend | None:
return None
def transform_session_created_event(
self,
model: str,

10
litellm/llms/bedrock/batches/handler.py
View file

@ -10,6 +10,8 @@ from litellm.types.llms.bedrock import AwsAuthParams, AwsSessionTag
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from botocore.config import Config
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
# AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses.
@ -31,6 +33,12 @@ _BEDROCK_MIJ_STATUS_TO_OPENAI: Final = {
_CANCEL_IDEMPOTENT_STATUSES: Final = frozenset({"cancelling", "cancelled", "completed", "failed", "expired"})
def _sigv4_config() -> "Config":
from botocore.config import Config
return Config(signature_version="v4")
def _extract_region_from_bedrock_arn(arn: str) -> str | None:
"""ARN shape: ``arn:aws:bedrock:<region>:<account>:<type>/<id>``"""
try:
@ -150,6 +158,7 @@ class BedrockBatchesHandler:
aws_access_key_id=creds.access_key,
aws_secret_access_key=creds.secret_key,
aws_session_token=creds.token,
config=_sigv4_config(),
)
def job_status() -> "LiteLLMBatch":
@ -309,6 +318,7 @@ class BedrockBatchesHandler:
aws_access_key_id=creds.access_key,
aws_secret_access_key=creds.secret_key,
aws_session_token=creds.token,
config=_sigv4_config(),
)
if logging_obj is not None:

11
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -108,6 +108,7 @@ BEDROCK_COMPUTER_USE_TOOLS: Final = [
"bash_",
"text_editor_",
]
BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS: Final = 16
# Beta header patterns that are not supported by Bedrock Converse API
# These will be filtered out to prevent errors
@ -378,6 +379,10 @@ class AmazonConverseConfig(BaseConfig):
def _is_openai_gpt_reasoning_model(model: str) -> bool:
return re.search(r"openai\.gpt-\d", model) is not None
@staticmethod
def _requires_min_max_tokens(model: str) -> bool:
return re.search(r"openai\.gpt-\d|xai\.grok-", model) is not None
def _is_nova_2_model(self, model: str) -> bool:
"""
Check if the model is a Nova 2 model that supports reasoningConfig.
@ -1000,7 +1005,11 @@ class AmazonConverseConfig(BaseConfig):
is_thinking_enabled=is_thinking_enabled,
)
if param == "max_tokens" or param == "max_completion_tokens":
optional_params["maxTokens"] = value
optional_params["maxTokens"] = (
max(value, BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS)
if isinstance(value, int) and self._requires_min_max_tokens(model)
else value
)
if param == "stream":
optional_params["stream"] = value
if param == "stop":

6
litellm/llms/custom_httpx/http_handler.py
View file

@ -150,7 +150,11 @@ def get_default_headers() -> dict:
if user_agent is not None:
return {"User-Agent": user_agent}
return {"User-Agent": f"litellm/{version}"}
return {"User-Agent": default_user_agent()}
def default_user_agent() -> str:
return f"litellm/{version}"
# Initialize headers (User-Agent)

50
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -2223,6 +2223,7 @@ class BaseLLMHTTPHandler:
# Prepare headers
kwargs = kwargs or {}
kwargs_for_agentic: Final = self._agentic_hook_kwargs(kwargs=kwargs, api_key=api_key, api_base=api_base)
provider_specific_header: Final = cast(
litellm.types.utils.ProviderSpecificHeader | Sequence[litellm.types.utils.ProviderSpecificHeader] | None,
kwargs.get("provider_specific_header", None),
@ -2410,7 +2411,7 @@ class BaseLLMHTTPHandler:
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
kwargs=kwargs_for_agentic,
hold_back=bool(held_back_tool_names),
server_fulfilled_tool_names=held_back_tool_names,
)
@ -2433,8 +2434,7 @@ class BaseLLMHTTPHandler:
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
kwargs=kwargs,
kwargs=kwargs_for_agentic,
)
async def _finalize_anthropic_messages_response(
@ -2447,14 +2447,8 @@ class BaseLLMHTTPHandler:
anthropic_messages_optional_request_params: dict,
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str,
api_key: str | None,
kwargs: dict,
kwargs: dict[str, object],
) -> AnthropicMessagesResponse | AsyncIterator:
# Inject api_key into kwargs so follow-up calls in agentic hooks can
# authenticate. api_key is a named param here (not in kwargs), so
# _prepare_followup_kwargs would miss it otherwise.
kwargs_for_agentic: Final = {**kwargs, "api_key": api_key} if api_key else kwargs
# Call agentic completion hooks (non-streaming path only)
final_response: Final = await self._call_agentic_completion_hooks(
response=initial_response,
model=model,
@ -2464,7 +2458,7 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
stream=False,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs_for_agentic,
kwargs=kwargs,
)
return self._maybe_wrap_in_fake_stream(
@ -5312,6 +5306,15 @@ class BaseLLMHTTPHandler:
fingerprints: Final = list(kwargs.get("_agentic_loop_fingerprints", []) or [])
return depth, max_loops, fingerprints
@staticmethod
def _agentic_hook_kwargs(
kwargs: Mapping[str, object], api_key: str | None, api_base: str | None
) -> dict[str, object]:
"""``api_key`` and ``api_base`` are named parameters of ``anthropic_messages`` rather than kwargs, so the
follow-up call an agentic hook makes only reaches the same deployment if they are re-added here."""
deployment_params: Final = {"api_key": api_key, "api_base": api_base}
return {**kwargs, **{key: value for key, value in deployment_params.items() if value}}
@staticmethod
def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool:
"""
@ -5906,7 +5909,15 @@ class BaseLLMHTTPHandler:
callback.__class__.__name__,
plan.stop_reason,
)
return self._maybe_wrap_in_fake_stream(response, logging_obj, api_surface)
return self._maybe_wrap_in_fake_stream(
await callback.async_post_agentic_loop_response_hook(
response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls),
plan=plan,
kwargs=kwargs_with_provider,
),
logging_obj,
api_surface,
)
if not plan.run_agentic_loop:
continue
@ -6279,7 +6290,12 @@ class BaseLLMHTTPHandler:
ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE
backend_ws: Final = await self._open_realtime_backend_ws(websockets, url, headers, ssl_context)
provider_backend: Final = await provider_config.open_backend(url, headers)
backend_ws: Final = (
provider_backend
if provider_backend is not None
else await self._open_realtime_backend_ws(websockets, url, headers, ssl_context)
)
async with backend_ws:
_request_data: Final[dict[str, object]] = {}
if litellm_metadata:
@ -6591,7 +6607,7 @@ class BaseLLMHTTPHandler:
first_message: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
**kwargs: Any,
):
) -> Exception | None:
"""
Handles Responses API WebSocket mode.
@ -6625,7 +6641,7 @@ class BaseLLMHTTPHandler:
**kwargs,
)
await handler.run()
return
return None
import websockets
from websockets.asyncio.client import ClientConnection
@ -6744,9 +6760,10 @@ class BaseLLMHTTPHandler:
output_guardrail_callbacks=_ws_output_guardrail_callbacks,
quota_callbacks=_ws_quota_callbacks,
authorized_model=model,
custom_llm_provider=custom_llm_provider,
request_defaults=request_defaults,
)
await streaming.bidirectional_forward()
return await streaming.bidirectional_forward()
except websockets.exceptions.InvalidStatusCode as e:
verbose_logger.exception("Error connecting to responses WS backend: %s", e)
@ -6760,6 +6777,7 @@ class BaseLLMHTTPHandler:
pass
else:
raise Exception(f"Unexpected error while closing WebSocket: {close_error}")
return None
def image_edit_handler(
self,

42
litellm/llms/fireworks_ai/cache_pricing.py Normal file
View file

@ -0,0 +1,42 @@
from typing import (
Final,
cast, # noqa: TID251 # the derived entry is a dict copy of a ReadOnly TypedDict; no cast-free way to retype it
)
from litellm.constants import FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO
from litellm.types.utils import ModelInfo
def _as_rate(value: object) -> float | None:
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
return None
try:
return float(value)
except ValueError:
return None
def with_default_cache_read_rate(model_info: ModelInfo) -> ModelInfo:
input_rate: Final = _as_rate(model_info.get("input_cost_per_token"))
if model_info.get("cache_read_input_token_cost") is not None or input_rate is None:
return model_info
cache_read_rate: Final = input_rate * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO
off_peak: Final = model_info.get("off_peak_pricing")
if off_peak is None or "cache_read_input_token_cost" in off_peak:
return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": cache_read_rate})
off_peak_input_rate: Final = _as_rate(off_peak.get("input_cost_per_token"))
return cast(
ModelInfo,
{
**model_info,
"cache_read_input_token_cost": cache_read_rate,
"off_peak_pricing": {
**off_peak,
"cache_read_input_token_cost": (
off_peak_input_rate * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO
if off_peak_input_rate is not None
else cache_read_rate
),
},
},
)

29
litellm/llms/fireworks_ai/cost_calculator.py
View file

@ -3,10 +3,7 @@ For calculating cost of fireworks ai serverless inference models.
"""
from datetime import datetime
from typing import (
Final,
cast, # noqa: TID251 # the fallback entry is a dict copy of a ReadOnly TypedDict; no cast-free way to retype it
)
from typing import Final
from litellm.constants import (
FIREWORKS_AI_4_B,
@ -67,28 +64,6 @@ def _resolve_model_info(model: str) -> ModelInfo:
return get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
def _with_cache_read_fallback(model_info: ModelInfo) -> ModelInfo:
"""Entries without a cache-read rate keep the previous calculator's input-rate fallback for cached
reads (LIT-7845 tracks the documented discount); the shared map is never mutated, so a copy carries it."""
input_rate: Final = model_info.get("input_cost_per_token")
if model_info.get("cache_read_input_token_cost") is not None or input_rate is None:
return model_info
off_peak: Final = model_info.get("off_peak_pricing")
if off_peak is None or "cache_read_input_token_cost" in off_peak:
return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": input_rate})
return cast(
ModelInfo,
{
**model_info,
"cache_read_input_token_cost": input_rate,
"off_peak_pricing": {
**off_peak,
"cache_read_input_token_cost": off_peak.get("input_cost_per_token", input_rate),
},
},
)
def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens,
@ -102,7 +77,7 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
"""
model_info: Final = _with_cache_read_fallback(_resolve_model_info(model))
model_info: Final = _resolve_model_info(model)
return generic_cost_per_token(
model=model,
usage=usage,

248
litellm/llms/meta/realtime/transformation.py
View file

@ -1,36 +1,41 @@
import asyncio
import base64
import binascii
import json
import math
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal
from typing import Final
from urllib.parse import urlparse, urlunparse
from pydantic import JsonValue, TypeAdapter, ValidationError
from pydantic import JsonValue
from litellm import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.realtime.transcription_protocol import (
RealtimeTranscriptionProtocolError,
TranscriptionSessionUpdate,
completed_event,
decode_pcm16_append,
delta_event,
duration_usage,
error_event,
json_object,
parse_transcription_session_update,
speech_event,
transcription_session,
transcription_session_created_event,
)
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.meta import MuseAudioEncoding, MuseHandshake, MuseMode, MuseSampleRate
from litellm.types.llms.openai import (
OpenAIRealtimeErrorEvent,
OpenAIRealtimeEvents,
OpenAIRealtimeInputAudioBufferSpeechEvent,
OpenAIRealtimeInputAudioTranscriptionCompleted,
OpenAIRealtimeInputAudioTranscriptionDelta,
OpenAIRealtimeServerVadTurnDetection,
OpenAIRealtimeTranscriptionSession,
OpenAIRealtimeTranscriptionSessionCreated,
OpenAIRealtimeTranscriptionSettings,
)
from litellm.types.realtime import (
RealtimeInputAudioTranscriptionDurationUsage,
RealtimeInputAudioTranscriptionUsage,
RealtimeResponseTransformInput,
RealtimeResponseTypedDict,
@ -98,17 +103,13 @@ _LANGUAGE_CODES: Final = MappingProxyType(
"zh": "Mandarin Chinese",
}
)
_SUPPORTED_TRANSCRIPTION_KEYS: Final = frozenset(("model", "language"))
_MAX_AUDIO_BACKLOG_SECONDS: Final = 4
_PACKET_MS: Final = 80
_END_STREAM: Final = '{"type":"endStream"}'
_PROVIDER_ERROR_MESSAGE: Final = "Meta Muse realtime transcription failed"
_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_EMPTY_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({})
_SERVER_VAD: Final[OpenAIRealtimeServerVadTurnDetection] = {"type": "server_vad"}
class MuseProtocolError(ValueError):
class MuseProtocolError(RealtimeTranscriptionProtocolError):
pass
@ -150,26 +151,13 @@ class MuseSessionConfig:
return biased
def openai_session(self, session_id: str) -> OpenAIRealtimeTranscriptionSession:
session: Final[OpenAIRealtimeTranscriptionSession] = {
"id": session_id,
"object": "realtime.transcription_session",
"type": "transcription",
"audio": {
"input": {
"format": {"type": "audio/pcm", "rate": self.sample_rate},
"transcription": self._transcription_settings(),
"turn_detection": None if self.mode == "PUSH_TO_TALK" else _SERVER_VAD,
}
},
}
return session
def _transcription_settings(self) -> OpenAIRealtimeTranscriptionSettings:
base: Final[OpenAIRealtimeTranscriptionSettings] = {"model": self.model}
if not self.language_bias:
return base
localized: Final[OpenAIRealtimeTranscriptionSettings] = {**base, "language": self.language_bias[0]}
return localized
return transcription_session(
session_id=session_id,
model=self.model,
sample_rate=self.sample_rate,
language=self.language_bias[0] if self.language_bias else None,
server_vad=self.mode != "PUSH_TO_TALK",
)
_DEFAULT_SESSION_CONFIG: Final = MuseSessionConfig(
@ -177,40 +165,10 @@ _DEFAULT_SESSION_CONFIG: Final = MuseSessionConfig(
)
def _json_object(payload: str) -> Mapping[str, JsonValue]:
try:
value: Final = _JSON_ADAPTER.validate_json(payload)
except ValidationError:
raise MuseProtocolError("invalid JSON object") from None
if not isinstance(value, dict):
raise MuseProtocolError("message must be a JSON object")
return value
def _mapping(value: JsonValue | None, name: str) -> Mapping[str, JsonValue]:
if value is None:
return _EMPTY_OBJECT
if not isinstance(value, dict):
raise MuseProtocolError(f"{name} must be an object")
return value
def _string(value: JsonValue | None, name: str) -> str | None:
if value is None:
return None
if not isinstance(value, str):
raise MuseProtocolError(f"{name} must be a string")
return value
def _normalize_model(model: str) -> str:
return model.removeprefix("meta/").strip()
def _event_id() -> str:
return f"event_{uuid.uuid4().hex}"
def normalize_language(language: str) -> str:
value: Final = language.strip()
if not value:
@ -254,138 +212,55 @@ def build_muse_realtime_url(api_base: str | None) -> str:
return urlunparse((scheme, netloc, "/v1/asr/realtime", "", "", ""))
def _parse_sample_rate(session: Mapping[str, JsonValue]) -> MuseSampleRate:
beta_format: Final = session.get("input_audio_format")
audio: Final = _mapping(session.get("audio"), "session.audio")
audio_input: Final = _mapping(audio.get("input"), "session.audio.input")
ga_format: Final = audio_input.get("format")
if beta_format is not None and ga_format is not None:
raise MuseProtocolError("input audio format must use either beta or GA layout")
if beta_format is not None:
if beta_format != "pcm16":
def _parse_sample_rate(update: TranscriptionSessionUpdate) -> MuseSampleRate:
audio_format: Final = update.audio_format
if audio_format is None:
return 24_000
if audio_format.layout == "beta":
if audio_format.encoding != "pcm16":
raise MuseProtocolError("Muse Voice requires pcm16 input audio")
return 24_000
if ga_format is None:
return 24_000
if isinstance(ga_format, str):
if ga_format != "pcm16":
raise MuseProtocolError("Muse Voice requires audio/pcm input audio")
return 24_000
format_mapping: Final = _mapping(ga_format, "session.audio.input.format")
if format_mapping.get("type") != "audio/pcm":
if not audio_format.is_pcm16:
raise MuseProtocolError("Muse Voice requires audio/pcm input audio")
channels: Final = format_mapping.get("channels", 1)
if isinstance(channels, bool) or channels != 1:
if audio_format.channels not in (None, 1):
raise MuseProtocolError("Muse Voice requires mono input audio")
rate: Final = format_mapping.get("rate", 24_000)
if isinstance(rate, bool) or not isinstance(rate, int) or rate not in SUPPORTED_SAMPLE_RATES:
rate: Final = 24_000 if audio_format.rate is None else audio_format.rate
if rate not in SUPPORTED_SAMPLE_RATES:
raise MuseProtocolError("Muse Voice supports PCM16 at 16000 Hz or 24000 Hz")
return 16_000 if rate == 16_000 else 24_000
def _parse_mode(session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue]) -> MuseMode:
turn_detection_present: Final = "turn_detection" in session or "turn_detection" in audio_input
turn_detection: Final = session.get("turn_detection", audio_input.get("turn_detection"))
if turn_detection_present and turn_detection is None:
def _parse_mode(update: TranscriptionSessionUpdate) -> MuseMode:
if update.turn_detection_disabled:
return "PUSH_TO_TALK"
if turn_detection is None:
return "ENDPOINTING"
turn_detection_mapping: Final = _mapping(turn_detection, "turn_detection")
if turn_detection_mapping.get("type") not in (None, "server_vad"):
if update.turn_detection_type not in (None, "server_vad"):
raise MuseProtocolError("Muse Voice supports server_vad turn detection or null")
return "ENDPOINTING"
def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig:
message: Final = _json_object(payload)
if message.get("type") not in ("session.update", "transcription_session.update"):
raise MuseProtocolError("expected session.update")
session: Final = _mapping(message.get("session"), "session")
if not session:
raise MuseProtocolError("session.update requires a session object")
if session.get("type") not in (None, "transcription", "realtime"):
update: Final = parse_transcription_session_update(payload, MuseProtocolError)
if update.session_type not in (None, "transcription", "realtime"):
raise MuseProtocolError("Muse Voice supports transcription sessions only")
audio: Final = _mapping(session.get("audio"), "session.audio")
audio_input: Final = _mapping(audio.get("input"), "session.audio.input")
beta_transcription: Final = session.get("input_audio_transcription")
ga_transcription: Final = audio_input.get("transcription")
if beta_transcription is not None and ga_transcription is not None:
raise MuseProtocolError("input transcription must use either beta or GA layout")
transcription: Final = _mapping(
beta_transcription if beta_transcription is not None else ga_transcription,
"input audio transcription",
)
unsupported: Final = tuple(sorted(key for key in transcription if key not in _SUPPORTED_TRANSCRIPTION_KEYS))
if unsupported:
verbose_logger.warning("Meta realtime: dropping unsupported transcription settings %s", unsupported)
requested_model: Final = _string(transcription.get("model"), "transcription model")
if update.unsupported_transcription_keys:
verbose_logger.warning(
"Meta realtime: dropping unsupported transcription settings %s", update.unsupported_transcription_keys
)
normalized_model: Final = _normalize_model(expected_model)
if normalized_model != MUSE_MODEL:
raise MuseProtocolError("unsupported Meta realtime model")
if requested_model is not None and _normalize_model(requested_model) != normalized_model:
if update.model is not None and _normalize_model(update.model) != normalized_model:
raise MuseProtocolError("realtime session model cannot be changed")
language: Final = _string(transcription.get("language"), "language")
return MuseSessionConfig(
model=normalized_model,
mode=_parse_mode(session, audio_input),
sample_rate=_parse_sample_rate(session),
language_bias=() if language is None else (normalize_language(language),),
mode=_parse_mode(update),
sample_rate=_parse_sample_rate(update),
language_bias=() if update.language is None else (normalize_language(update.language),),
)
def session_created_event(config: MuseSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated:
event: Final[OpenAIRealtimeTranscriptionSessionCreated] = {
"type": "session.created",
"event_id": _event_id(),
"session": config.openai_session(session_id),
}
return event
def error_event(message: str) -> OpenAIRealtimeErrorEvent:
event: Final[OpenAIRealtimeErrorEvent] = {
"type": "error",
"error": {"type": "server_error", "message": message},
}
return event
def _speech_event(
event_type: Literal["input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"], item_id: str
) -> OpenAIRealtimeInputAudioBufferSpeechEvent:
event: Final[OpenAIRealtimeInputAudioBufferSpeechEvent] = {
"type": event_type,
"event_id": _event_id(),
"item_id": item_id,
}
return event
def _delta_event(item_id: str, delta: str) -> OpenAIRealtimeInputAudioTranscriptionDelta:
event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = {
"type": "conversation.item.input_audio_transcription.delta",
"event_id": _event_id(),
"item_id": item_id,
"content_index": 0,
"delta": delta,
}
return event
def _completed_event(
item_id: str, transcript: str, usage: RealtimeInputAudioTranscriptionUsage | None
) -> OpenAIRealtimeInputAudioTranscriptionCompleted:
event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = {
"type": "conversation.item.input_audio_transcription.completed",
"event_id": _event_id(),
"item_id": item_id,
"content_index": 0,
"transcript": transcript,
}
if usage is None:
return event
billed: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = {**event, "usage": usage}
return billed
return transcription_session_created_event(config.openai_session(session_id))
def _required_turn_id(message: Mapping[str, JsonValue], event: str) -> str:
@ -424,18 +299,18 @@ class _TurnState:
has_content: Final = self.latest_partial is not None or self.final_text is not None
if (self.started or has_content) and not self.start_emitted:
self.start_emitted = True
yield _speech_event("input_audio_buffer.speech_started", self.item_id)
yield speech_event("input_audio_buffer.speech_started", self.item_id)
if self.latest_partial is not None and self.final_text is None:
delta: Final = _new_suffix(self.emitted_partial, self.latest_partial)
if delta:
self.emitted_partial = self.latest_partial
yield _delta_event(self.item_id, delta)
yield delta_event(self.item_id, delta)
if self.stopped and not self.stopped_emitted:
self.stopped_emitted = True
yield _speech_event("input_audio_buffer.speech_stopped", self.item_id)
yield speech_event("input_audio_buffer.speech_stopped", self.item_id)
if self.final_text is not None and self.stopped_emitted and not self.completed_emitted:
self.completed_emitted = True
yield _completed_event(self.item_id, self.final_text, take_usage())
yield completed_event(self.item_id, self.final_text, take_usage())
class MuseEventTransformer:
@ -467,8 +342,7 @@ class MuseEventTransformer:
if seconds <= 0:
return None
self._unbilled_seconds = 0.0
usage: Final[RealtimeInputAudioTranscriptionDurationUsage] = {"type": "duration", "seconds": seconds}
return usage
return duration_usage(seconds)
def _apply_turn_event(self, event_type: JsonValue | None, message: Mapping[str, JsonValue]) -> _TurnState | None:
match event_type:
@ -612,7 +486,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig):
model: str,
session_configuration_request: str | None = None,
) -> tuple[str | bytes, ...]:
request: Final = _json_object(message)
request: Final = json_object(message, MuseProtocolError)
event_type: Final = request.get("type")
if event_type in ("session.update", "transcription_session.update"):
return self._configure(message, model)
@ -664,7 +538,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig):
return result
def _backend_events(self, payload: str) -> tuple[OpenAIRealtimeEvents, ...]:
frame: Final = _json_object(payload)
frame: Final = json_object(payload, MuseProtocolError)
session_id: Final = frame.get("sessionId")
if session_id is None:
return self._transformer.transform(frame)
@ -686,17 +560,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig):
def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]:
config: Final = self._require_config()
encoded: Final = request.get("audio")
if not isinstance(encoded, str):
raise MuseProtocolError("Audio must be a base64 string")
if len(encoded) > config.max_encoded_append_bytes:
raise MuseProtocolError("Audio append exceeds the four-second backlog limit")
try:
audio: Final = base64.b64decode(encoded, validate=True)
except (binascii.Error, ValueError):
raise MuseProtocolError("Audio must be valid base64") from None
if len(audio) % 2:
raise MuseProtocolError("PCM16 audio must contain complete samples")
audio: Final = decode_pcm16_append(request.get("audio"), config.max_encoded_append_bytes, MuseProtocolError)
buffered: Final = self._pending_audio + audio
packet_end: Final = len(buffered) - len(buffered) % config.packet_bytes
self._pending_audio = buffered[packet_end:]

418
litellm/llms/vertex_ai/audio_transcription/realtime_backend.py Normal file
View file

@ -0,0 +1,418 @@
import asyncio
import time
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
from dataclasses import dataclass
from datetime import timedelta
from types import MappingProxyType, TracebackType
from typing import TYPE_CHECKING, Final, Literal, Protocol
from pydantic import TypeAdapter
from typing_extensions import Self, assert_never
from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK
from websockets.frames import Close
from litellm import verbose_logger
from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import SpeechStreamingTarget
from litellm.types.llms.vertex_ai_speech_to_text import (
VertexSpeechStreamingCommand,
VertexSpeechStreamingCommandUnion,
VertexSpeechStreamingConfigure,
VertexSpeechStreamingConfigured,
VertexSpeechStreamingDiscardTurn,
VertexSpeechStreamingFinishTurn,
VertexSpeechStreamingResponse,
VertexSpeechStreamingResult,
VertexSpeechStreamingTurnDiscarded,
VertexSpeechStreamingTurnFinished,
)
if TYPE_CHECKING:
from google.cloud.speech_v2.types import (
StreamingRecognitionConfig,
StreamingRecognizeRequest,
StreamingRecognizeResponse,
)
SPEECH_SDK_INSTALL_HINT: Final = (
"google-cloud-speech is not installed. Install with `pip install 'litellm[stt-vertex-chirp]'`."
)
STREAM_FAILURE_CLOSE_CODE: Final = 1011
STREAM_ROTATION_SECONDS: Final = 240.0
STREAM_ROTATION_DEADLINE_SECONDS: Final = 280.0
REQUEST_QUEUE_SIZE: Final = 64
OUTBOX_SIZE: Final = 256
_LINK_QUEUE_SIZE: Final = 64
_CLOSE_REASON_MAX_CHARS: Final = 120
_CONFIGURED_EVENT: Final = VertexSpeechStreamingConfigured().model_dump_json()
_TURN_FINISHED_EVENT: Final = VertexSpeechStreamingTurnFinished().model_dump_json()
_COMMAND_ADAPTER: Final = TypeAdapter[VertexSpeechStreamingCommandUnion](VertexSpeechStreamingCommand)
_TIMEDELTA_ADAPTER: Final = TypeAdapter(timedelta)
_SPEECH_EVENTS: Final[MappingProxyType[str, Literal["begin", "end"]]] = MappingProxyType(
{
"SPEECH_ACTIVITY_BEGIN": "begin",
"SPEECH_ACTIVITY_END": "end",
"END_OF_SINGLE_UTTERANCE": "end",
}
)
class ClosableTransport(Protocol):
def close(self) -> Awaitable[None]: ...
class SpeechStreamingClient(Protocol):
def streaming_recognize(
self, requests: "AsyncIterator[StreamingRecognizeRequest] | None" = None
) -> "Awaitable[AsyncIterable[StreamingRecognizeResponse]]": ...
@property
def transport(self) -> ClosableTransport: ...
@dataclass(frozen=True, slots=True)
class _StreamFailure:
reason: str
@dataclass(frozen=True, slots=True)
class _Closed:
pass
@dataclass(frozen=True, slots=True)
class _TurnResult:
turn: int
event: str
@dataclass(frozen=True, slots=True)
class _TurnDiscarded:
turn: int
@dataclass(frozen=True, slots=True)
class _TurnDiscardedEvent:
turn: int
event: str
_OutboxItem = str | _TurnResult | _TurnDiscardedEvent | _StreamFailure | _Closed
def open_speech_client(target: SpeechStreamingTarget, access_token: str) -> SpeechStreamingClient:
try:
from google.api_core.client_options import ClientOptions
from google.cloud.speech_v2 import SpeechAsyncClient
from google.oauth2.credentials import Credentials
except ImportError as e:
raise ImportError(SPEECH_SDK_INSTALL_HINT) from e
return SpeechAsyncClient(
credentials=Credentials(token=access_token),
transport="grpc_asyncio",
client_options=ClientOptions(api_endpoint=target.api_endpoint),
)
def _streaming_config(command: VertexSpeechStreamingConfigure) -> "StreamingRecognitionConfig":
from google.cloud.speech_v2.types import (
ExplicitDecodingConfig,
RecognitionConfig,
StreamingRecognitionConfig,
StreamingRecognitionFeatures,
)
return StreamingRecognitionConfig(
config=RecognitionConfig(
explicit_decoding_config=ExplicitDecodingConfig(
encoding=ExplicitDecodingConfig.AudioEncoding.LINEAR16,
sample_rate_hertz=command.sample_rate_hertz,
audio_channel_count=1,
),
model=command.model,
language_codes=command.language_codes,
),
streaming_features=StreamingRecognitionFeatures(interim_results=True, enable_voice_activity_events=True),
)
def _response_event(response: "StreamingRecognizeResponse", billed_seconds: float) -> str:
return VertexSpeechStreamingResponse(
speech_event=_SPEECH_EVENTS.get(response.speech_event_type.name, "none"),
results=tuple(
VertexSpeechStreamingResult(
transcript=result.alternatives[0].transcript if result.alternatives else "",
is_final=result.is_final,
)
for result in response.results
),
billed_seconds=billed_seconds,
).model_dump_json()
def _billed_seconds(response: "StreamingRecognizeResponse") -> float:
return _TIMEDELTA_ADAPTER.validate_python(response.metadata.total_billed_duration).total_seconds()
def _normal_closure() -> ConnectionClosedOK:
return ConnectionClosedOK(rcvd=Close(1000, ""), sent=None)
class _RecognizeStream:
def __init__(
self,
*,
client: SpeechStreamingClient,
request_type: "type[StreamingRecognizeRequest]",
first_request: "StreamingRecognizeRequest",
opened_at: float,
turn: int,
) -> None:
self._client: Final = client
self._request_type: Final = request_type
self.opened_at: Final = opened_at
self.turn: Final = turn
self._requests: Final[asyncio.Queue[StreamingRecognizeRequest | None]] = asyncio.Queue(
maxsize=REQUEST_QUEUE_SIZE
)
self._requests.put_nowait(first_request)
self.speech_active: bool = False
self.billed_seconds: float = 0.0
self._cancelled: bool = False
self._closed: bool = False
self._task: asyncio.Task[None] | None = None
async def send_audio(self, audio: bytes) -> None:
await self._requests.put(self._request_type(audio=audio))
async def half_close(self) -> None:
await self._requests.put(None)
def cancel(self) -> None:
self._cancelled = True
if self._task is not None:
self._task.cancel()
async def close(self) -> None:
if self._closed:
return
self._closed = True
await self._client.transport.close()
async def relay(self, outbox: asyncio.Queue[_OutboxItem], billed_before: float) -> float:
if self._cancelled:
await self.close()
return 0.0
task: Final = asyncio.create_task(self._forward(outbox, billed_before))
self._task = task
try:
await asyncio.wait((task,))
except asyncio.CancelledError:
task.cancel()
await asyncio.wait((task,))
raise
finally:
await self.close()
return self.billed_seconds
async def _forward(self, outbox: asyncio.Queue[_OutboxItem], billed_before: float) -> None:
try:
responses: Final = await self._client.streaming_recognize(self._drain())
async for response in responses:
self._note(response)
await outbox.put(
_TurnResult(turn=self.turn, event=_response_event(response, billed_before + self.billed_seconds))
)
except Exception as e: # noqa: BLE001 # task boundary: a swallowed failure would hang the client session
verbose_logger.warning("Google Speech-to-Text streaming failed: %s", e)
await outbox.put(_StreamFailure(reason=f"Google Speech-to-Text streaming failed: {e}"))
def _note(self, response: "StreamingRecognizeResponse") -> None:
activity: Final = _SPEECH_EVENTS.get(response.speech_event_type.name)
if activity is not None:
self.speech_active = activity == "begin"
self.billed_seconds = max(self.billed_seconds, _billed_seconds(response))
async def _drain(self) -> "AsyncIterator[StreamingRecognizeRequest]":
while (request := await self._requests.get()) is not None:
yield request
_Link = _RecognizeStream | str | _TurnDiscarded
class SpeechStreamingBackend:
def __init__(
self,
target: SpeechStreamingTarget,
*,
client_factory: Callable[[SpeechStreamingTarget, str], SpeechStreamingClient] = open_speech_client,
clock: Callable[[], float] = time.monotonic,
rotation_seconds: float = STREAM_ROTATION_SECONDS,
rotation_deadline_seconds: float = STREAM_ROTATION_DEADLINE_SECONDS,
) -> None:
self._target: Final = target
self._client_factory: Final = client_factory
self._clock: Final = clock
self._rotation_seconds: Final = rotation_seconds
self._rotation_deadline_seconds: Final = rotation_deadline_seconds
self._outbox: Final[asyncio.Queue[_OutboxItem]] = asyncio.Queue(maxsize=OUTBOX_SIZE)
self._links: Final[asyncio.Queue[_Link]] = asyncio.Queue(maxsize=_LINK_QUEUE_SIZE)
self._pump: asyncio.Task[None] | None = None
self._config: StreamingRecognitionConfig | None = None
self._turn: tuple[_RecognizeStream, ...] = ()
self._turn_index: int = 0
self._discarded_turns: frozenset[int] = frozenset()
self._billed_before: float = 0.0
self._closed: bool = False
async def __aenter__(self) -> Self:
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
await self.close()
async def send(self, message: str | bytes) -> None:
if self._closed:
raise _normal_closure()
if isinstance(message, bytes):
await self._send_audio(message)
return
command: Final = _COMMAND_ADAPTER.validate_json(message)
match command:
case VertexSpeechStreamingConfigure():
self._config = _streaming_config(command)
await self._link(_CONFIGURED_EVENT)
case VertexSpeechStreamingFinishTurn():
await self._finish_turn()
case VertexSpeechStreamingDiscardTurn():
await self._discard_turn()
case _:
assert_never(command)
async def recv(self, decode: bool | None = None) -> str | bytes:
while not (self._closed and self._outbox.empty()):
if (event := self._deliverable(await self._outbox.get())) is not None:
return event
raise _normal_closure()
def _deliverable(self, item: _OutboxItem) -> str | None:
match item:
case _StreamFailure():
raise ConnectionClosedError(
rcvd=Close(STREAM_FAILURE_CLOSE_CODE, item.reason[:_CLOSE_REASON_MAX_CHARS]), sent=None
)
case _Closed():
raise _normal_closure()
case _TurnResult():
return None if item.turn in self._discarded_turns else item.event
case _TurnDiscardedEvent():
self._discarded_turns -= {item.turn}
return item.event
case str():
return item
case _:
assert_never(item)
async def close(self) -> None:
if self._closed:
return
self._closed = True
self._turn = ()
pump: Final = self._pump
if pump is not None:
pump.cancel()
await asyncio.wait((pump,))
await self._close_unrelayed_streams()
if not self._outbox.full():
self._outbox.put_nowait(_Closed())
async def _close_unrelayed_streams(self) -> None:
unrelayed: Final = tuple(self._links.get_nowait() for _ in range(self._links.qsize()))
for link in unrelayed:
if isinstance(link, _RecognizeStream):
await link.close()
async def _link(self, item: _Link) -> None:
if self._pump is None:
self._pump = asyncio.create_task(self._pump_links())
await self._links.put(item)
async def _pump_links(self) -> None:
while True:
await self._relay(await self._links.get())
async def _relay(self, link: _Link) -> None:
match link:
case str():
await self._outbox.put(link)
case _RecognizeStream():
self._billed_before += await link.relay(self._outbox, self._billed_before)
case _TurnDiscarded():
await self._outbox.put(
_TurnDiscardedEvent(
turn=link.turn,
event=VertexSpeechStreamingTurnDiscarded(billed_seconds=self._billed_before).model_dump_json(),
)
)
case _:
assert_never(link)
async def _send_audio(self, audio: bytes) -> None:
stream: Final = await self._turn_stream()
await stream.send_audio(audio)
async def _turn_stream(self) -> _RecognizeStream:
current: Final = self._turn[-1] if self._turn else None
if current is not None and not self._expired(current):
return current
if current is not None:
await current.half_close()
stream: Final = await self._open_stream()
self._turn = (*self._turn, stream)
return stream
def _expired(self, stream: _RecognizeStream) -> bool:
elapsed: Final = self._clock() - stream.opened_at
if elapsed >= self._rotation_deadline_seconds:
return True
return elapsed >= self._rotation_seconds and not stream.speech_active
async def _open_stream(self) -> _RecognizeStream:
from google.cloud.speech_v2.types import StreamingRecognizeRequest
config: Final = self._config
if config is None:
raise RuntimeError("audio was sent before the Speech-to-Text stream was configured")
access_token: Final = await self._target.resolve_access_token()
stream: Final = _RecognizeStream(
client=self._client_factory(self._target, access_token),
request_type=StreamingRecognizeRequest,
first_request=StreamingRecognizeRequest(recognizer=self._target.recognizer, streaming_config=config),
opened_at=self._clock(),
turn=self._turn_index,
)
await self._link(stream)
return stream
async def _finish_turn(self) -> None:
turn: Final = self._turn
self._turn = ()
self._turn_index += 1
if turn:
await turn[-1].half_close()
await self._link(_TURN_FINISHED_EVENT)
async def _discard_turn(self) -> None:
streams: Final = self._turn
turn: Final = self._turn_index
self._turn = ()
self._discarded_turns |= {turn}
self._turn_index += 1
for stream in streams:
stream.cancel()
await self._link(_TurnDiscarded(turn=turn))

446
litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py Normal file
View file

@ -0,0 +1,446 @@
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass, replace
from typing import Final
from pydantic import JsonValue, TypeAdapter
from typing_extensions import assert_never
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.audio_utils.utils import normalize_transcription_language_to_bcp47
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.realtime.transcription_protocol import (
RealtimeTranscriptionProtocolError,
TranscriptionAudioFormat,
TranscriptionSessionUpdate,
completed_event,
decode_pcm16_append,
delta_event,
duration_usage,
json_object,
parse_transcription_session_update,
speech_event,
transcription_session,
transcription_session_created_event,
)
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend
from litellm.llms.vertex_ai.audio_transcription.transformation import (
AUTO_LANGUAGE_CODE,
DEFAULT_SPEECH_TO_TEXT_LOCATION,
speech_to_text_host,
validate_vertex_transcription_location,
validate_vertex_transcription_project_id,
)
from litellm.types.llms.openai import (
OpenAIRealtimeEvents,
OpenAIRealtimeTranscriptionSession,
OpenAIRealtimeTranscriptionSessionCreated,
)
from litellm.types.llms.vertex_ai_speech_to_text import (
VertexSpeechStreamingConfigure,
VertexSpeechStreamingConfigured,
VertexSpeechStreamingDiscardTurn,
VertexSpeechStreamingEvent,
VertexSpeechStreamingEventUnion,
VertexSpeechStreamingFinishTurn,
VertexSpeechStreamingResponse,
VertexSpeechStreamingTurnDiscarded,
VertexSpeechStreamingTurnFinished,
)
from litellm.types.realtime import (
RealtimeInputAudioTranscriptionUsage,
RealtimeResponseTransformInput,
RealtimeResponseTypedDict,
)
DEFAULT_SAMPLE_RATE_HERTZ: Final = 24_000
MIN_SAMPLE_RATE_HERTZ: Final = 8_000
MAX_SAMPLE_RATE_HERTZ: Final = 48_000
MAX_AUDIO_MESSAGE_BYTES: Final = 25_000
_SPEECH_TO_TEXT_ENDPOINTS: Final = frozenset({"/v1/audio/transcriptions", "/v1/realtime"})
_VERTEX_MODEL_PREFIX: Final = "vertex_ai/"
_STREAMING_EVENT_ADAPTER: Final = TypeAdapter[VertexSpeechStreamingEventUnion](VertexSpeechStreamingEvent)
_FINISH_TURN_COMMAND: Final = VertexSpeechStreamingFinishTurn().model_dump_json()
_DISCARD_TURN_COMMAND: Final = VertexSpeechStreamingDiscardTurn().model_dump_json()
class ChirpProtocolError(RealtimeTranscriptionProtocolError):
pass
@dataclass(frozen=True, slots=True)
class SpeechStreamingTarget:
api_endpoint: str
recognizer: str
resolve_access_token: Callable[[], Awaitable[str]]
@dataclass(frozen=True, slots=True)
class ChirpSessionConfig:
model: str
language: str | None
sample_rate: int
server_vad: bool
def openai_session(self, session_id: str) -> OpenAIRealtimeTranscriptionSession:
return transcription_session(
session_id=session_id,
model=self.model,
sample_rate=self.sample_rate,
language=self.language,
server_vad=self.server_vad,
)
def configure_command(self) -> str:
return VertexSpeechStreamingConfigure(
model=self.model,
language_codes=(AUTO_LANGUAGE_CODE,) if self.language is None else (self.language,),
sample_rate_hertz=self.sample_rate,
).model_dump_json()
def is_vertex_speech_to_text_model(model: str) -> bool:
try:
info: Final = litellm.get_model_info(
model=normalize_speech_to_text_model(model), custom_llm_provider="vertex_ai"
)
except Exception: # noqa: BLE001 # get_model_info raises for unmapped models, which are not Speech-to-Text models
return False
if info.get("mode") != "audio_transcription":
return False
return _SPEECH_TO_TEXT_ENDPOINTS <= frozenset(info.get("supported_endpoints") or ())
def normalize_speech_to_text_model(model: str) -> str:
return model.removeprefix(_VERTEX_MODEL_PREFIX)
def default_session_config(model: str) -> ChirpSessionConfig:
return ChirpSessionConfig(
model=normalize_speech_to_text_model(model),
language=None,
sample_rate=DEFAULT_SAMPLE_RATE_HERTZ,
server_vad=True,
)
def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessionConfig:
update: Final = parse_transcription_session_update(payload, ChirpProtocolError)
if update.session_type not in (None, "transcription", "realtime"):
raise ChirpProtocolError("Speech-to-Text streaming supports transcription sessions only")
if update.unsupported_transcription_keys:
verbose_logger.debug(
"Speech-to-Text streaming: ignoring unsupported transcription settings %s",
update.unsupported_transcription_keys,
)
model: Final = normalize_speech_to_text_model(expected_model)
if update.model is not None and normalize_speech_to_text_model(update.model) != model:
raise ChirpProtocolError("realtime session model cannot be changed")
return ChirpSessionConfig(
model=model,
language=None if update.language is None else normalize_transcription_language_to_bcp47(update.language),
sample_rate=_parse_sample_rate(update.audio_format),
server_vad=_parse_server_vad(update),
)
def _parse_sample_rate(audio_format: TranscriptionAudioFormat | None) -> int:
if audio_format is None:
return DEFAULT_SAMPLE_RATE_HERTZ
if not audio_format.is_pcm16:
raise ChirpProtocolError("Speech-to-Text streaming requires pcm16 input audio")
if audio_format.channels not in (None, 1):
raise ChirpProtocolError("Speech-to-Text streaming requires mono input audio")
rate: Final = DEFAULT_SAMPLE_RATE_HERTZ if audio_format.rate is None else audio_format.rate
if not MIN_SAMPLE_RATE_HERTZ <= rate <= MAX_SAMPLE_RATE_HERTZ:
raise ChirpProtocolError(
f"Speech-to-Text streaming supports sample rates from {MIN_SAMPLE_RATE_HERTZ} Hz"
f" to {MAX_SAMPLE_RATE_HERTZ} Hz"
)
return rate
def _parse_server_vad(update: TranscriptionSessionUpdate) -> bool:
if update.turn_detection_disabled:
return False
if update.turn_detection_type not in (None, "server_vad"):
raise ChirpProtocolError("Speech-to-Text streaming supports server_vad turn detection or null")
return True
def session_created_event(config: ChirpSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated:
return transcription_session_created_event(config.openai_session(session_id))
def _normalize_word(word: str) -> str:
return "".join(char for char in word if char.isalnum()).casefold()
def new_words(previous: str, current: str) -> str:
previous_words: Final = previous.split()
current_words: Final = current.split()
common: Final = next(
(
index
for index, (old, new) in enumerate(zip(previous_words, current_words, strict=False))
if _normalize_word(old) != _normalize_word(new)
),
min(len(previous_words), len(current_words)),
)
appended: Final = " ".join(current_words[common:])
if not appended:
return ""
return f" {appended}" if common else appended
def _join_transcript(committed: str, tail: str) -> str:
return " ".join(part for part in (committed, tail) if part)
@dataclass(frozen=True, slots=True)
class _Turn:
item_id: str
committed: str = ""
preview: str = ""
started_emitted: bool = False
stopped_emitted: bool = False
class ChirpEventTransformer:
def __init__(self, *, new_item_id: Callable[[], str] = lambda: f"item_{uuid.uuid4().hex}") -> None:
self._new_item_id: Final = new_item_id
self._config: ChirpSessionConfig | None = None
self._session_id: str | None = None
self._turn: _Turn | None = None
self._billed_seconds: float = 0.0
self._reported_seconds: float = 0.0
def configure(self, config: ChirpSessionConfig, session_id: str) -> None:
self._config = config
self._session_id = session_id
def take_unbilled_usage(self) -> RealtimeInputAudioTranscriptionUsage | None:
unreported: Final = self._billed_seconds - self._reported_seconds
if unreported <= 0:
return None
self._reported_seconds = self._billed_seconds
return duration_usage(unreported)
def transform(self, frame: VertexSpeechStreamingEventUnion) -> tuple[OpenAIRealtimeEvents, ...]:
match frame:
case VertexSpeechStreamingConfigured():
return (session_created_event(self._require_config(), self._require_session_id()),)
case VertexSpeechStreamingResponse():
return self._response(frame)
case VertexSpeechStreamingTurnFinished():
return self._finish_turn()
case VertexSpeechStreamingTurnDiscarded():
self._billed_seconds = max(self._billed_seconds, frame.billed_seconds)
self._turn = None
return ()
case _:
assert_never(frame)
def _response(self, frame: VertexSpeechStreamingResponse) -> tuple[OpenAIRealtimeEvents, ...]:
self._billed_seconds = max(self._billed_seconds, frame.billed_seconds)
interim: Final = " ".join(
result.transcript.strip() for result in frame.results if not result.is_final and result.transcript.strip()
)
finals: Final = tuple(
result.transcript.strip() for result in frame.results if result.is_final and result.transcript.strip()
)
begin_events: Final = self._begin() if frame.speech_event == "begin" else ()
final_events: Final = tuple(event for final in finals for event in self._final(final))
interim_events: Final = self._hypothesis(interim) if interim else ()
end_events: Final = self._stop() if frame.speech_event == "end" else ()
return (*begin_events, *final_events, *interim_events, *end_events)
def _begin(self) -> tuple[OpenAIRealtimeEvents, ...]:
turn: Final = self._require_turn()
if turn.started_emitted or not self._require_config().server_vad:
return ()
self._turn = replace(turn, started_emitted=True)
return (speech_event("input_audio_buffer.speech_started", turn.item_id),)
def _stop(self) -> tuple[OpenAIRealtimeEvents, ...]:
turn: Final = self._turn
if turn is None or turn.stopped_emitted or not self._require_config().server_vad:
return ()
self._turn = replace(turn, stopped_emitted=True)
return (speech_event("input_audio_buffer.speech_stopped", turn.item_id),)
def _hypothesis(self, text: str) -> tuple[OpenAIRealtimeEvents, ...]:
begin_events: Final = self._begin()
turn: Final = self._require_turn()
hypothesis: Final = _join_transcript(turn.committed, text)
delta: Final = new_words(turn.preview, hypothesis)
self._turn = replace(turn, preview=hypothesis)
return (*begin_events, delta_event(turn.item_id, delta)) if delta else begin_events
def _final(self, text: str) -> tuple[OpenAIRealtimeEvents, ...]:
begin_events: Final = self._begin()
turn: Final = self._require_turn()
committed: Final = _join_transcript(turn.committed, text)
delta: Final = new_words(turn.preview, committed)
self._turn = replace(turn, committed=committed, preview=committed)
delta_events: Final[tuple[OpenAIRealtimeEvents, ...]] = (delta_event(turn.item_id, delta),) if delta else ()
if not self._require_config().server_vad:
return (*begin_events, *delta_events)
return (*begin_events, *delta_events, *self._complete())
def _finish_turn(self) -> tuple[OpenAIRealtimeEvents, ...]:
if self._turn is None:
return ()
return self._complete()
def _complete(self) -> tuple[OpenAIRealtimeEvents, ...]:
turn: Final = self._require_turn()
stop_events: Final = self._stop()
transcript: Final = turn.committed or turn.preview
self._turn = None
return (*stop_events, completed_event(turn.item_id, transcript, self.take_unbilled_usage()))
def _require_turn(self) -> _Turn:
if self._turn is None:
self._turn = _Turn(item_id=self._new_item_id())
return self._turn
def _require_config(self) -> ChirpSessionConfig:
if self._config is None:
raise ChirpProtocolError("session.update must configure the session before the backend responds")
return self._config
def _require_session_id(self) -> str:
if self._session_id is None:
raise ChirpProtocolError("session.update must configure the session before the backend responds")
return self._session_id
def _default_backend_factory(target: SpeechStreamingTarget) -> RealtimeBackend:
from litellm.llms.vertex_ai.audio_transcription.realtime_backend import SpeechStreamingBackend
return SpeechStreamingBackend(target)
class VertexChirpRealtimeConfig(BaseRealtimeConfig):
def __init__(
self,
*,
resolve_access_token: Callable[[], Awaitable[str]],
project: str,
location: str | None,
backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory,
) -> None:
self._resolve_access_token: Final = resolve_access_token
self._project: Final = validate_vertex_transcription_project_id(project)
self._location: Final = validate_vertex_transcription_location(location, DEFAULT_SPEECH_TO_TEXT_LOCATION)
self._backend_factory: Final = backend_factory
self._transformer: Final = ChirpEventTransformer()
self._config: ChirpSessionConfig | None = None
self._session_id: str | None = None
def validate_environment(
self,
headers: dict[str, str], # mutable-ok: BaseRealtimeConfig contract
model: str,
api_key: str | None = None,
) -> dict[str, str]: # mutable-ok: BaseRealtimeConfig contract
return headers
def get_complete_url(self, api_base: str | None, model: str, api_key: str | None = None) -> str:
if not is_vertex_speech_to_text_model(model):
raise ValueError(f"Unsupported Speech-to-Text streaming model: {model}")
return _api_endpoint(api_base) if api_base else speech_to_text_host(self._location)
async def open_backend(self, url: str, headers: Mapping[str, str]) -> RealtimeBackend | None:
return self._backend_factory(
SpeechStreamingTarget(
api_endpoint=url,
recognizer=f"projects/{self._project}/locations/{self._location}/recognizers/_",
resolve_access_token=self._resolve_access_token,
)
)
def is_setup_message(self, msg_obj: Mapping[str, object]) -> bool:
return msg_obj.get("kind") == "configure"
def transform_session_created_event(
self,
model: str,
logging_session_id: str,
session_configuration_request: str | None = None,
) -> OpenAIRealtimeTranscriptionSessionCreated:
self._session_id = logging_session_id
return session_created_event(default_session_config(model), logging_session_id)
def transform_realtime_request(
self,
message: str,
model: str,
session_configuration_request: str | None = None,
) -> tuple[str | bytes, ...]:
request: Final = json_object(message, ChirpProtocolError)
event_type: Final = request.get("type")
if event_type in ("session.update", "transcription_session.update"):
return self._configure(message, model)
if event_type == "input_audio_buffer.append":
return self._append_audio(request)
if event_type in ("input_audio_buffer.commit", "input_audio_buffer.end"):
self._require_config()
return (_FINISH_TURN_COMMAND,)
if event_type == "input_audio_buffer.clear":
self._require_config()
return (_DISCARD_TURN_COMMAND,)
verbose_logger.debug("Speech-to-Text streaming: dropping unsupported client event %s", event_type)
return ()
def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None:
return self._transformer.take_unbilled_usage()
def transform_realtime_response(
self,
message: str | bytes,
model: str,
logging_obj: LiteLLMLoggingObj,
realtime_response_transform_input: RealtimeResponseTransformInput,
) -> RealtimeResponseTypedDict:
frame: Final = _STREAMING_EVENT_ADAPTER.validate_json(message)
events: Final = list(self._transformer.transform(frame)) # mutable-ok: response field is a list
result: Final[RealtimeResponseTypedDict] = {
"response": events,
"current_output_item_id": realtime_response_transform_input.get("current_output_item_id"),
"current_response_id": realtime_response_transform_input.get("current_response_id"),
"current_delta_chunks": realtime_response_transform_input.get("current_delta_chunks"),
"current_conversation_id": realtime_response_transform_input.get("current_conversation_id"),
"current_item_chunks": realtime_response_transform_input.get("current_item_chunks"),
"current_delta_type": realtime_response_transform_input.get("current_delta_type"),
"session_configuration_request": realtime_response_transform_input.get("session_configuration_request"),
}
return result
def _configure(self, message: str, model: str) -> tuple[str, ...]:
if self._config is not None:
verbose_logger.debug("Speech-to-Text streaming: ignoring session.update after the stream was configured")
return ()
config: Final = parse_chirp_session_update(message, model)
self._config = config
self._transformer.configure(config, self._session_id or f"sess_{uuid.uuid4().hex}")
return (config.configure_command(),)
def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]:
self._require_config()
audio: Final = decode_pcm16_append(request.get("audio"), error=ChirpProtocolError)
return tuple(
audio[start : start + MAX_AUDIO_MESSAGE_BYTES] for start in range(0, len(audio), MAX_AUDIO_MESSAGE_BYTES)
)
def _require_config(self) -> ChirpSessionConfig:
if self._config is None:
raise ChirpProtocolError("session.update must configure the session before audio is sent")
return self._config
def _api_endpoint(api_base: str) -> str:
without_scheme: Final = api_base.split("://", 1)[-1]
return without_scheme.split("/", 1)[0]

7
litellm/llms/vertex_ai/audio_transcription/transformation.py
View file

@ -42,6 +42,10 @@ def validate_vertex_transcription_location(location: str | None, default_locatio
raise VertexAIError(status_code=400, message=str(e)) from e
def speech_to_text_host(location: str) -> str:
return "speech.googleapis.com" if location == "global" else f"{location}-speech.googleapis.com"
def validate_vertex_transcription_project_id(project_id: str) -> str:
if not project_id or ".." in project_id or any(c in project_id for c in _URL_UNSAFE_PROJECT_CHARS):
raise VertexAIError(status_code=400, message=f"Invalid vertex_project format: {project_id!r}")
@ -122,8 +126,7 @@ class VertexAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase)
project_id: Final = validate_vertex_transcription_project_id(
self.safe_get_vertex_ai_project(litellm_params) or self._resolve_project_id_from_credentials(litellm_params)
)
host: Final = "speech.googleapis.com" if location == "global" else f"{location}-speech.googleapis.com"
base_url: Final = (api_base or f"https://{host}").rstrip("/")
base_url: Final = (api_base or f"https://{speech_to_text_host(location)}").rstrip("/")
return f"{base_url}/v2/projects/{project_id}/locations/{location}/recognizers/_:recognize"
def _resolve_project_id_from_credentials(self, litellm_params: dict) -> str:

23
litellm/llms/vertex_ai/realtime/transformation.py
View file

@ -12,10 +12,16 @@ Auth: OAuth2 Bearer token (not an API key).
"""
import json
from collections.abc import Awaitable, Callable
from typing import Final
from litellm import verbose_logger
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import (
VertexChirpRealtimeConfig,
is_vertex_speech_to_text_model,
)
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
class VertexAIRealtimeConfig(GeminiRealtimeConfig):
@ -232,3 +238,20 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
return []
return super().transform_realtime_request(message, model, session_configuration_request)
def vertex_realtime_config(
model: str,
*,
access_token: str,
resolve_access_token: Callable[[], Awaitable[str]],
project: str,
location: str | None,
) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig:
if is_vertex_speech_to_text_model(model):
return VertexChirpRealtimeConfig(resolve_access_token=resolve_access_token, project=project, location=location)
return VertexAIRealtimeConfig(
access_token=access_token,
project=project,
location=VertexBase.get_vertex_region(vertex_region=location, model=model),
)

3
litellm/llms/xai/audio_transcription/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .transformation import XAIAudioTranscriptionConfig
__all__ = ["XAIAudioTranscriptionConfig"]

207
litellm/llms/xai/audio_transcription/transformation.py Normal file
View file

@ -0,0 +1,207 @@
"""
Translates from OpenAI's `/v1/audio/transcriptions` to xAI's `/v1/stt`
"""
from collections.abc import Mapping, Sequence
from typing import Final
from httpx import Headers, Response
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIAudioTranscriptionOptionalParams,
)
from litellm.types.utils import FileTypes, TranscriptionResponse
from ...base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
BaseAudioTranscriptionConfig,
)
from ..common_utils import XAIModelInfo
class XAIAudioTranscriptionError(BaseLLMException):
pass
class _XAISttWord(BaseModel):
model_config = ConfigDict(extra="allow")
text: str = ""
start: float = 0.0
end: float = 0.0
speaker: int | None = None
class _XAISttResponse(BaseModel):
model_config = ConfigDict(extra="allow")
text: str = ""
language: str = "unknown"
duration: float | None = None
words: tuple[_XAISttWord, ...] | None = None
_OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...])
_STRING_OBJECT_DICT: Final = TypeAdapter(dict[str, object])
def _serialize_form_value(
value: object,
) -> str | list[str]: # mutable-ok: httpx multipart data takes list values for repeated form fields
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, (list, tuple)):
return [str(item) for item in _OBJECT_TUPLE.validate_python(value)]
return str(value)
class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
@property
def custom_llm_provider(self) -> str:
return litellm.LlmProviders.XAI.value
@property
def has_native_transcription_endpoint(self) -> bool:
return True
def get_supported_openai_params(
self, model: str
) -> list[OpenAIAudioTranscriptionOptionalParams]: # mutable-ok: base class signature returns list
return ["language"]
def map_openai_params(
self,
non_default_params: Mapping[str, object],
optional_params: Mapping[str, object],
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: base class signature returns dict
supported_params: Final = self.get_supported_openai_params(model)
return {
**optional_params,
**{k: v for k, v in non_default_params.items() if k in supported_params},
}
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | Headers, # mutable-ok: base class signature takes dict
) -> BaseLLMException:
return XAIAudioTranscriptionError(message=error_message, status_code=status_code, headers=headers)
def transform_audio_transcription_request(
self,
model: str,
audio_file: FileTypes,
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> AudioTranscriptionRequestData:
processed_audio: Final = process_audio_file(audio_file)
extra_body: Final = optional_params.get("extra_body")
flat_params: Final[Mapping[str, object]] = {
**(_STRING_OBJECT_DICT.validate_python(extra_body) if isinstance(extra_body, Mapping) else {}),
**{k: v for k, v in optional_params.items() if k != "extra_body"},
}
excluded_params: Final = frozenset({"model", "OPENAI_TRANSCRIPTION_PARAMS", "extra_body"})
form_data: Final[
dict[str, str | list[str]]
] = { # mutable-ok: AudioTranscriptionRequestData.data requires dict and httpx needs list values
"model": model,
**{
k: _serialize_form_value(v)
for k, v in flat_params.items()
if v is not None and k not in excluded_params
},
}
files: Final = {
"file": (
processed_audio.filename,
processed_audio.file_content,
processed_audio.content_type,
)
}
return AudioTranscriptionRequestData(data=form_data, files=files)
def transform_audio_transcription_response(
self,
raw_response: Response,
) -> TranscriptionResponse:
if raw_response.status_code >= 400:
raise self.get_error_class(
error_message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
try:
payload: Final = _XAISttResponse.model_validate_json(raw_response.content)
except ValidationError as e:
raise XAIAudioTranscriptionError(
message=f"Error parsing xAI response: {e}",
status_code=raw_response.status_code,
headers=dict(raw_response.headers),
)
response: Final = TranscriptionResponse(text=payload.text)
response["task"] = "transcribe"
response["language"] = payload.language
if payload.duration is not None:
response["duration"] = payload.duration
if payload.words is not None:
response["words"] = [
{
"word": word.text,
"start": word.start,
"end": word.end,
**({"speaker": word.speaker} if word.speaker is not None else {}),
}
for word in payload.words
]
hidden_params: Final[dict[str, object]] = dict(
payload.model_dump(mode="json")
) # mutable-ok: TranscriptionResponse._hidden_params is a dict
if payload.duration is not None:
hidden_params["audio_transcription_duration"] = payload.duration
response._hidden_params = hidden_params # pyright: ignore[reportPrivateUsage] # TranscriptionResponse exposes no public hidden-params setter
return response
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
stream: bool | None = None,
) -> str:
base: Final = (XAIModelInfo.get_api_base(api_base) or "").rstrip("/")
normalized: Final = base.removesuffix("/v1")
return f"{normalized}/v1/stt"
def validate_environment(
self,
headers: dict[str, object], # mutable-ok: base class signature takes and returns dict
model: str,
messages: Sequence[AllMessageValues],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, object]: # mutable-ok: base class signature returns dict
resolved_key: Final = XAIModelInfo.get_api_key(api_key)
if resolved_key is None:
raise ValueError("xAI API key is required. Set XAI_API_KEY environment variable.")
return {**headers, "Authorization": f"Bearer {resolved_key}"}

7
litellm/main.py
View file

@ -64,6 +64,7 @@ from litellm.constants import (
AZURE_OPENAI_AUDIO_PROVIDERS,
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS,
)
from litellm.exceptions import LiteLLMUnknownProvider
from litellm.integrations.custom_logger import CustomLogger
@ -7830,6 +7831,10 @@ def transcription(
provider=LlmProviders(custom_llm_provider),
)
uses_openai_transport: Final = custom_llm_provider in OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS and not (
provider_config is not None and provider_config.has_native_transcription_endpoint
)
if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None:
# azure configs
api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@ -7859,7 +7864,7 @@ def transcription(
litellm_params=litellm_params_dict,
custom_llm_provider=custom_llm_provider,
)
elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers):
elif uses_openai_transport:
api_base = (
api_base
or litellm.api_base

299
litellm/model_prices_and_context_window_backup.json
View file

@ -19244,7 +19244,20 @@
"supports_function_calling": true,
"supports_anthropic_thinking_payload": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_reasoning": true,
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text"
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
},
"databricks/databricks-gemini-2-5-pro": {
"cache_creation_input_token_cost": 1.24999e-06,
@ -19265,7 +19278,21 @@
"supports_function_calling": true,
"supports_anthropic_thinking_payload": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"deprecation_date": "2026-10-02",
"supports_reasoning": true,
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text"
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
},
"databricks/databricks-gemini-3-1-flash-lite": {
"cache_creation_input_token_cost": 3.1248e-07,
@ -19285,7 +19312,19 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text"
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
},
"databricks/databricks-gemini-3-1-flash-image": {
"litellm_provider": "databricks",
@ -19347,7 +19386,20 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_reasoning": true,
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text"
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
},
"databricks/databricks-gemini-3-flash": {
"cache_creation_input_token_cost": 6.2503e-07,
@ -19367,7 +19419,19 @@
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text"
],
"supports_vision": true,
"supports_audio_input": true,
"supports_video_input": true
},
"databricks/databricks-gemini-3-pro": {
"cache_creation_input_token_cost": 2.49998e-06,
@ -21433,7 +21497,11 @@
"mode": "chat",
"supports_tool_choice": true,
"supports_function_calling": true,
"supports_image_size": false
"supports_image_size": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_audio_input": true
},
"deepinfra/google/gemini-2.5-pro": {
"max_tokens": 1000000,
@ -21444,7 +21512,11 @@
"litellm_provider": "deepinfra",
"mode": "chat",
"supports_tool_choice": true,
"supports_function_calling": true
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_audio_input": true
},
"deepinfra/google/gemma-3-12b-it": {
"max_tokens": 131072,
@ -29340,26 +29412,6 @@
"supports_parallel_function_calling": true,
"supports_vision": true
},
"github_copilot/gemini-2.5-pro": {
"litellm_provider": "github_copilot",
"max_input_tokens": 128000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_vision": true
},
"github_copilot/gemini-3-pro-preview": {
"litellm_provider": "github_copilot",
"max_input_tokens": 128000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_vision": true
},
"github_copilot/gpt-3.5-turbo": {
"litellm_provider": "github_copilot",
"max_input_tokens": 16384,
@ -30014,17 +30066,6 @@
"output_cost_per_token": 8.8e-07,
"supports_function_calling": true
},
"gmi/google/gemini-3-pro-preview": {
"input_cost_per_token": 2e-06,
"litellm_provider": "gmi",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 1.2e-05,
"supports_function_calling": true,
"supports_vision": true
},
"gmi/google/gemini-3-flash-preview": {
"input_cost_per_token": 5e-07,
"litellm_provider": "gmi",
@ -30034,7 +30075,8 @@
"mode": "chat",
"output_cost_per_token": 3e-06,
"supports_function_calling": true,
"supports_vision": true
"supports_vision": true,
"supports_system_messages": true
},
"gmi/moonshotai/Kimi-K2-Thinking": {
"input_cost_per_token": 8e-07,
@ -40163,7 +40205,12 @@
"supports_response_schema": true,
"supports_vision": true,
"supports_native_streaming": true,
"supports_image_size": false
"supports_image_size": false,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_pdf_input": true,
"supports_audio_input": true,
"supports_video_input": true
},
"oci/google.gemini-2.5-pro": {
"input_cost_per_token": 1.25e-06,
@ -40177,7 +40224,12 @@
"supports_function_calling": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_native_streaming": true
"supports_native_streaming": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_pdf_input": true,
"supports_audio_input": true,
"supports_video_input": true
},
"oci/google.gemini-2.5-flash-lite": {
"input_cost_per_token": 7.5e-08,
@ -40192,7 +40244,12 @@
"supports_response_schema": true,
"supports_vision": true,
"supports_native_streaming": true,
"supports_image_size": false
"supports_image_size": false,
"supports_reasoning": false,
"supports_system_messages": true,
"supports_pdf_input": true,
"supports_audio_input": true,
"supports_video_input": true
},
"oci/cohere.command-a-vision": {
"input_cost_per_token": 1.56e-06,
@ -41435,7 +41492,7 @@
"max_tokens": 65535,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"supports_audio_output": true,
"supports_audio_output": false,
"supports_function_calling": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -41449,7 +41506,8 @@
"supports_audio_input": true,
"supports_pdf_input": true,
"supports_reasoning": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-2.5-pro": {
"cache_creation_input_token_cost": 3.75e-07,
@ -41462,7 +41520,7 @@
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 1e-05,
"supports_audio_output": true,
"supports_audio_output": false,
"supports_function_calling": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -41478,7 +41536,8 @@
"supports_audio_input": true,
"supports_pdf_input": true,
"supports_reasoning": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3-pro-preview": {
"cache_read_input_token_cost": 2e-07,
@ -41563,7 +41622,8 @@
"supports_url_context": true,
"supports_vision": true,
"supports_web_search": false,
"tpm": 800000
"tpm": 800000,
"supports_video_input": true
},
"openrouter/google/gemini-3.1-flash-lite-preview": {
"cache_creation_input_token_cost": 8.33333333333333e-08,
@ -41690,7 +41750,8 @@
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/gryphe/mythomax-l2-13b": {
"input_cost_per_token": 8e-08,
@ -44413,12 +44474,16 @@
"output_cost_per_token": 1.2e-05,
"litellm_provider": "replicate",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_function_calling": false,
"supports_parallel_function_calling": false,
"supports_vision": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_response_schema": true
"supports_tool_choice": false,
"supports_response_schema": false,
"input_cost_per_token_above_200k_tokens": 4e-06,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"supports_audio_input": true,
"supports_video_input": true
},
"replicate/anthropic/claude-4.5-sonnet": {
"input_cost_per_token": 3e-06,
@ -44487,17 +44552,19 @@
"supports_response_schema": true
},
"replicate/google/gemini-2.5-flash": {
"input_cost_per_token": 2.5e-06,
"input_cost_per_token": 3e-07,
"output_cost_per_token": 2.5e-06,
"litellm_provider": "replicate",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_function_calling": false,
"supports_parallel_function_calling": false,
"supports_vision": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_response_schema": true,
"supports_image_size": false
"supports_tool_choice": false,
"supports_response_schema": false,
"supports_image_size": false,
"supports_reasoning": true,
"supports_video_input": true
},
"replicate/openai/gpt-oss-120b": {
"input_cost_per_token": 1.8e-07,
@ -48076,10 +48143,15 @@
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true,
"supports_image_size": false
"supports_image_size": false,
"cache_read_input_token_cost": 3e-08,
"supports_reasoning": true,
"supports_pdf_input": true,
"supports_web_search": true,
"supports_prompt_caching": true
},
"vercel_ai_gateway/google/gemini-2.5-pro": {
"input_cost_per_token": 2.5e-06,
"input_cost_per_token": 1.25e-06,
"litellm_provider": "vercel_ai_gateway",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
@ -48089,7 +48161,15 @@
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
"supports_response_schema": true,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"output_cost_per_token_above_200k_tokens": 1.5e-05,
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_above_200k_tokens": 2.5e-07,
"supports_reasoning": true,
"supports_pdf_input": true,
"supports_web_search": true,
"supports_prompt_caching": true
},
"vercel_ai_gateway/google/gemini-embedding-001": {
"input_cost_per_token": 1.5e-07,
@ -48825,7 +48905,8 @@
"mode": "audio_transcription",
"source": "https://cloud.google.com/speech-to-text/pricing",
"supported_endpoints": [
"/v1/audio/transcriptions"
"/v1/audio/transcriptions",
"/v1/realtime"
]
},
"vertex_ai/claude-3-5-haiku": {
@ -62224,7 +62305,8 @@
"supports_response_schema": true,
"supports_reasoning": true,
"supports_vision": true,
"source": "https://deepinfra.com/pricing"
"source": "https://deepinfra.com/pricing",
"supports_audio_input": true
},
"deepinfra/XiaomiMiMo/MiMo-V2.5": {
"max_tokens": 262144,
@ -62524,7 +62606,8 @@
"supports_response_schema": true,
"supports_reasoning": true,
"supports_vision": true,
"source": "https://deepinfra.com/pricing"
"source": "https://deepinfra.com/pricing",
"supports_audio_input": true
},
"deepinfra/google/gemini-3.7-flash": {
"max_tokens": 1000000,
@ -62538,7 +62621,8 @@
"supports_response_schema": true,
"supports_reasoning": true,
"supports_vision": true,
"source": "https://deepinfra.com/pricing"
"source": "https://deepinfra.com/pricing",
"supports_audio_input": true
},
"deepinfra/inclusionAI/Ling-3.0-flash": {
"max_tokens": 131072,
@ -62970,7 +63054,8 @@
"supports_response_schema": true,
"supports_reasoning": true,
"supports_vision": true,
"source": "https://deepinfra.com/pricing"
"source": "https://deepinfra.com/pricing",
"supports_audio_input": true
},
"deepinfra/XiaomiMiMo/MiMo-V2.5-Pro": {
"max_tokens": 1048576,
@ -63375,6 +63460,34 @@
"video"
]
},
"xai/grok-voice-transcribe-1.0": {
"input_cost_per_second": 2.778e-05,
"litellm_provider": "xai",
"metadata": {
"calculation": "$0.10/3600 seconds = $0.00002778 per second",
"original_pricing_per_hour": 0.1
},
"mode": "audio_transcription",
"output_cost_per_second": 0.0,
"source": "https://docs.x.ai/developers/pricing",
"supported_endpoints": [
"/v1/audio/transcriptions"
]
},
"xai/grok-voice-transcribe-2.0": {
"input_cost_per_second": 2.778e-05,
"litellm_provider": "xai",
"metadata": {
"calculation": "$0.10/3600 seconds = $0.00002778 per second",
"original_pricing_per_hour": 0.1
},
"mode": "audio_transcription",
"output_cost_per_second": 0.0,
"source": "https://docs.x.ai/developers/pricing",
"supported_endpoints": [
"/v1/audio/transcriptions"
]
},
"low/1024-x-1024/grok-imagine-image-2.0": {
"input_cost_per_image": 0.04,
"litellm_provider": "xai",
@ -65365,7 +65478,8 @@
"deprecation_date": "2026-10-20",
"input_cost_per_audio_token": 3e-07,
"supports_prompt_caching": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.5-flash": {
"cache_creation_input_token_cost": 8.33333333333333e-08,
@ -65388,7 +65502,8 @@
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_audio_token": 3e-06,
"supports_prompt_caching": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.5-flash-lite": {
"cache_creation_input_token_cost": 8.33333333333333e-08,
@ -65411,7 +65526,8 @@
"cache_read_input_token_cost": 3e-08,
"input_cost_per_audio_token": 3e-07,
"supports_prompt_caching": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.6-flash": {
"cache_creation_input_token_cost": 4.16666666666667e-08,
@ -65434,7 +65550,8 @@
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 7.5e-07,
"supports_prompt_caching": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.7-flash": {
"cache_creation_input_token_cost": 4.16666666666667e-08,
@ -65457,7 +65574,8 @@
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 7.5e-07,
"supports_prompt_caching": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.8-flash": {
"cache_creation_input_token_cost": 4.16666666666667e-08,
@ -65480,7 +65598,8 @@
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 7.5e-07,
"supports_prompt_caching": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/openai/gpt-4o-mini": {
"input_cost_per_token": 1.5e-07,
@ -67108,7 +67227,8 @@
"supports_pdf_input": true,
"supports_audio_input": true,
"supports_prompt_caching": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/qwen/qwen3-max-thinking": {
"input_cost_per_token": 7.8e-07,
@ -71974,7 +72094,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-2.5-flash:batch": {
"cache_read_input_audio_token_cost": 1e-07,
@ -71997,7 +72118,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-2.5-pro:batch": {
"cache_read_input_audio_token_cost": 1.25e-07,
@ -72023,7 +72145,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3-flash-preview:batch": {
"input_cost_per_audio_token": 5e-07,
@ -72043,7 +72166,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.1-flash-lite:batch": {
"cache_read_input_audio_token_cost": 2.5e-08,
@ -72065,7 +72189,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.1-pro-preview:batch": {
"input_cost_per_audio_token": 1e-06,
@ -72087,7 +72212,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.5-flash-lite:batch": {
"cache_read_input_audio_token_cost": 1.5e-08,
@ -72109,7 +72235,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.5-flash:batch": {
"cache_read_input_audio_token_cost": 1.5e-07,
@ -72131,7 +72258,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.6-flash:batch": {
"cache_creation_input_token_cost": 4.16666666666667e-08,
@ -72154,7 +72282,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.7-flash:batch": {
"cache_creation_input_token_cost": 4.16666666666667e-08,
@ -72177,7 +72306,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/google/gemini-3.8-flash:batch": {
"cache_creation_input_token_cost": 4.16666666666667e-08,
@ -72200,7 +72330,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"supports_video_input": true
},
"openrouter/ibm-granite/granite-4.0-h-micro": {
"input_cost_per_token": 1.7e-08,

10
litellm/proxy/_experimental/mcp_server/AGENTS.md
View file

@ -1,6 +1,6 @@
# Experimental MCP Server Change Guidelines
Read @../../../../CLAUDE.md and @CLAUDE.md before changing this package.
Read @../../../../AGENTS.md before changing this package.
This directory owns the proxy-hosted MCP server implementation. Keep changes
inside the module that owns the behavior, and only reach outside this package
@ -14,13 +14,13 @@ Respect the current package boundaries:
```text
litellm/proxy/_experimental/mcp_server/
AGENTS.md
CLAUDE.md
server.py # ASGI/MCP route handling, sessions, tool calls [PR7: 7-arm only — move BYOK/OAuth pre-fetch into resolver]
mcp_server_manager.py # upstream server registry, clients, tool routing [PR7: _create_mcp_client swaps resolve_mcp_auth -> resolve_credentials]
auth/
user_api_key_auth_mcp.py # LiteLLM admission auth and MCP request headers
token_exchange.py # OAuth token exchange handling [unchanged; V1TokenExchangeAdapter delegates here]
litellm_auth_handler.py # authenticated-user adapter for MCP sessions
client_allowlist.py # gateway-level client application allowlist (mcp_allowed_clients); leaf module, no litellm.proxy imports
outbound_credentials/ # NEW — typed upstream-credential resolution (resolve_credentials + arms)
__init__.py # public surface: resolve_credentials, the configs, CredError
result.py # Ok | Error union (pure stdlib)
@ -67,8 +67,10 @@ module materially harder to understand.
auth, SSE, streamable HTTP, and stdio as separate flows. Do not collapse them
behind a single generic branch unless tests prove every mode still behaves
correctly.
- Be especially careful with legacy `delegate_auth_to_upstream: true`. The local
`CLAUDE.md` explains its admitted replacement and public discovery contract.
- Be especially careful with legacy `delegate_auth_to_upstream: true`. `auth_type: oauth2`
with `delegate_auth_to_upstream: true` is deprecated: LiteLLM admission is required
for matching MCP routes. Use `auth_type: oauth_delegate` for client-forwarded OAuth.
OAuth discovery endpoints stay public so clients can start the RFC 9728 flow.
- Keep database-backed fields in sync across migrations, typed models under
`litellm/types/mcp.py` or `litellm/types/mcp_server/`, config loading, this
package, and dashboard state when the field is user-visible.

1
litellm/proxy/_experimental/mcp_server/CLAUDE.md
View file

@ -1 +0,0 @@
MCP note: **`auth_type: oauth2` with `delegate_auth_to_upstream: true` is deprecated** - LiteLLM admission is required for matching MCP routes. Use `auth_type: oauth_delegate` for client-forwarded OAuth. OAuth discovery endpoints stay public so clients can start the RFC 9728 flow

14
litellm/proxy/_experimental/mcp_server/bridge_token_flow.py
View file

@ -262,7 +262,12 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No
return loaded if isinstance(loaded, str) else None
async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResolutionFailure":
UserRowSource = Literal["cache", "database"]
async def load_active_user_by_id(
user_id: str, source: UserRowSource = "cache"
) -> "LiteLLM_UserTable | _KeyResolutionFailure":
"""Load a live litellm user by id, returning the record when the user is active or a precise
failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a
user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a
@ -273,7 +278,11 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol
``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object``
catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look
identical, the original error surviving only as ``__context__``), so the outage check walks the cause
chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault."""
chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.
``source="database"`` reads the row from the database, never the cache, so the credential mint refuses
a user that a writer deactivated or deleted without evicting the cached row, and it leaves the fresh
row in the cache for the requests the credential makes next. Every other caller keeps the cache read,
so introspection, which a resource server may call per request, stays off the database."""
from litellm.proxy._types import (
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
)
@ -296,6 +305,7 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
check_db_only=source == "database",
)
except (ProxyException, HTTPException):
return "no_active_key"

171
litellm/proxy/_experimental/mcp_server/client_allowlist.py Normal file
View file

@ -0,0 +1,171 @@
"""
Gateway-level allowlist of MCP client applications (``general_settings.mcp_allowed_clients``).
Each entry pairs an admin-chosen ``alias`` (shown in the dashboard and logs) with the ``value`` that
identifies the client. Only the value is compared, exactly and case-sensitively.
A caller that authenticated with a JWT is identified by the claim named in
``litellm_jwtauth.mcp_client_id_jwt_field``, a value asserted by the identity provider.
Every other caller is identified by the header named in ``general_settings.mcp_client_id_header``,
which the client picks itself, so that source is a policy control rather than a security boundary.
While the allowlist is set, a caller with no usable identity source is rejected.
"""
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
from litellm.types.mcp import MCPAllowedClient
MCP_ALLOWED_CLIENTS_SETTING: Final = "mcp_allowed_clients"
MCP_CLIENT_ID_HEADER_SETTING: Final = "mcp_client_id_header"
MCP_CLIENT_ID_JWT_FIELD_SETTING: Final = "mcp_client_id_jwt_field"
_JWT_AUTH_SETTING: Final = "litellm_jwtauth"
_ALLOWED_CLIENTS_ADAPTER: Final[TypeAdapter[list[MCPAllowedClient]]] = TypeAdapter(list[MCPAllowedClient])
_OPTIONAL_NAME_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
_OPTIONAL_MAPPING_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter(dict[str, object] | None)
_NOBODY: Final[Mapping[str, str]] = MappingProxyType({})
class MCPClientForbiddenBody(TypedDict):
error: ReadOnly[Literal["Forbidden"]]
details: ReadOnly[str]
@dataclass(frozen=True, slots=True)
class MCPClientAllowlist:
"""``aliases_by_value`` maps each admitted identity value to the alias the admin gave it."""
aliases_by_value: Mapping[str, str]
jwt_field: str | None
header: str | None
@dataclass(frozen=True, slots=True)
class MCPClientIdentity:
client_id: str
source: Literal["jwt", "header"]
source_name: str
@property
def description(self) -> str:
return f"'{self.client_id}' (from {'JWT claim' if self.source == 'jwt' else 'header'} '{self.source_name}')"
@dataclass(frozen=True, slots=True)
class MCPClientRejection:
details: str
@property
def response_body(self) -> MCPClientForbiddenBody:
body: Final[MCPClientForbiddenBody] = {"error": "Forbidden", "details": self.details}
return body
def _unidentified_rejection(reason: str) -> MCPClientRejection:
return MCPClientRejection(
details=f"{reason} This gateway only admits client applications listed in {MCP_ALLOWED_CLIENTS_SETTING}."
)
def parse_allowed_mcp_clients(raw_setting: object) -> Mapping[str, str] | None:
"""Value-to-alias mapping; None when the setting is absent (not enforced). A malformed setting admits nobody."""
if raw_setting is None:
return None
try:
clients: Final = _ALLOWED_CLIENTS_ADAPTER.validate_python(raw_setting)
except ValidationError:
verbose_logger.warning(
"%s is not a list of {alias, value} entries (%r); rejecting every MCP client until it is fixed",
MCP_ALLOWED_CLIENTS_SETTING,
raw_setting,
)
return _NOBODY
return MappingProxyType({client.value: client.alias for client in clients})
def _parse_optional_name(setting_name: str, raw_setting: object) -> str | None:
try:
name: Final = _OPTIONAL_NAME_ADAPTER.validate_python(raw_setting)
except ValidationError:
verbose_logger.warning("%s is not a string (%r); ignoring it", setting_name, raw_setting)
return None
return name or None
def _jwt_field_from_general_settings(general_settings: Mapping[str, object]) -> str | None:
try:
jwt_auth: Final = _OPTIONAL_MAPPING_ADAPTER.validate_python(general_settings.get(_JWT_AUTH_SETTING))
except ValidationError:
return None
if jwt_auth is None:
return None
return _parse_optional_name(
f"{_JWT_AUTH_SETTING}.{MCP_CLIENT_ID_JWT_FIELD_SETTING}", jwt_auth.get(MCP_CLIENT_ID_JWT_FIELD_SETTING)
)
def load_mcp_client_allowlist(general_settings: Mapping[str, object]) -> MCPClientAllowlist | None:
"""None when ``mcp_allowed_clients`` is unset, which admits every client."""
allowed_clients: Final = parse_allowed_mcp_clients(general_settings.get(MCP_ALLOWED_CLIENTS_SETTING))
if allowed_clients is None:
return None
header: Final = _parse_optional_name(
MCP_CLIENT_ID_HEADER_SETTING, general_settings.get(MCP_CLIENT_ID_HEADER_SETTING)
)
return MCPClientAllowlist(
aliases_by_value=allowed_clients,
jwt_field=_jwt_field_from_general_settings(general_settings),
header=header.lower() if header is not None else None,
)
def resolve_mcp_client_identity(
allowlist: MCPClientAllowlist,
jwt_claims: Mapping[str, object] | None,
headers: Mapping[str, str],
) -> MCPClientIdentity | MCPClientRejection:
"""A JWT caller is identified by its configured claim alone, so a header can never override the IdP."""
if jwt_claims is not None and allowlist.jwt_field is not None:
claim: Final[object] = get_nested_value(data=jwt_claims, key_path=allowlist.jwt_field)
if isinstance(claim, str) and claim:
return MCPClientIdentity(client_id=claim, source="jwt", source_name=allowlist.jwt_field)
return _unidentified_rejection(
f"The JWT presented has no '{allowlist.jwt_field}' claim naming the client application."
)
if allowlist.header is None:
configured: Final = (
f"litellm_jwtauth.{MCP_CLIENT_ID_JWT_FIELD_SETTING} for JWT callers or {MCP_CLIENT_ID_HEADER_SETTING}"
)
return _unidentified_rejection(
f"No client identity source is configured for this request; set {configured} in general_settings."
)
header_value: Final = headers.get(allowlist.header)
if header_value:
return MCPClientIdentity(client_id=header_value, source="header", source_name=allowlist.header)
return _unidentified_rejection(f"The request has no '{allowlist.header}' header naming the client application.")
def check_mcp_client_allowed(
allowlist: MCPClientAllowlist | None,
jwt_claims: Mapping[str, object] | None,
headers: Mapping[str, str],
) -> MCPClientRejection | None:
if allowlist is None:
return None
identity: Final = resolve_mcp_client_identity(allowlist, jwt_claims, headers)
if isinstance(identity, MCPClientRejection):
return identity
alias: Final = allowlist.aliases_by_value.get(identity.client_id)
if alias is None:
return MCPClientRejection(
details=f"MCP client {identity.description} is not listed in this gateway's {MCP_ALLOWED_CLIENTS_SETTING}."
)
verbose_logger.debug("Admitted MCP client '%s' identified as %s", alias, identity.description)
return None

26
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -59,6 +59,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
register_aggregate_client,
relative_request_url,
revoke_refresh_token,
supported_grant_types,
)
from litellm.proxy._experimental.mcp_server.idp_token_exchange import (
exchange_idp_subject_token,
token_exchange_available,
)
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import (
RefreshOwnershipProven,
@ -1980,6 +1985,9 @@ async def token_endpoint(
refresh_token: str | None = Form(None),
scope: str | None = Form(None),
resource: str | None = Form(None),
subject_token: str | None = Form(None),
subject_token_type: str | None = Form(None),
requested_token_type: str | None = Form(None),
mcp_server_name: str | None = None,
):
"""
@ -2010,6 +2018,10 @@ async def token_endpoint(
cache=user_api_key_cache,
resource=resource,
mint_proxy_credential=mint_proxy_credential,
subject_token=subject_token,
subject_token_type=subject_token_type,
requested_token_type=requested_token_type,
exchange_subject_token=exchange_idp_subject_token,
)
lookup_name: Final = mcp_server_name or client_id
@ -2131,7 +2143,9 @@ async def introspect_endpoint(token: str = Form(...)) -> Response:
async def native_client_auth_discovery(request: Request) -> JSONResponse:
"""The versioned contract a native client (``lite login --pkce``, or a CLI in any other
language) reads to sign a user in through the browser and obtain a proxy credential."""
return JSONResponse(native_client_auth_contract(request), headers=TOKEN_NO_CACHE_HEADERS)
return JSONResponse(
native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS
)
# Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request
@ -2619,7 +2633,7 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict:
}
def _build_aggregate_authorization_server_response(request: Request) -> dict:
def _build_aggregate_authorization_server_response(request: Request, token_exchange_available: bool) -> dict:
"""RFC 8414 metadata for the gateway as the aggregate authorization server.
The issuer is ``{base}/mcp`` and must stay equal to the value the
@ -2638,7 +2652,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict:
"registration_endpoint": f"{request_base_url}/register",
"response_types_supported": ["code"],
"scopes_supported": [],
"grant_types_supported": ["authorization_code", "refresh_token"],
"grant_types_supported": supported_grant_types(token_exchange_available),
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["none", "client_secret_post"],
}
@ -2676,7 +2690,7 @@ async def oauth_authorization_server_aggregate(request: Request):
per-server row win here instead would serve an issuer of {base} against a resource that
advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.
"""
return _build_aggregate_authorization_server_response(request)
return _build_aggregate_authorization_server_response(request, token_exchange_available())
# Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name}
@ -2902,7 +2916,9 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
# advertises that), so this does not affect it. A request without redirect_uris is not
# a DCR request, so the legacy single-server-or-dummy fallback is kept for it.
if data.get("redirect_uris"):
return await register_aggregate_client(request=request, request_body=data)
return await register_aggregate_client(
request=request, request_body=data, token_exchange_available=token_exchange_available()
)
resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if resolved:
return await register_client_with_server(

175
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -51,7 +51,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
from fastapi import HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from pydantic import BaseModel, ConfigDict, Field, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
@ -187,6 +187,52 @@ class MintProxyCredential(Protocol):
) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ...
TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]:
"""The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed
only where the JWT auth that proves a subject token is on, backed by a database, and
licensed, so a client never selects a grant the gateway would then refuse."""
if token_exchange_available:
return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE)
return ("authorization_code", "refresh_token")
"""RFC 8693: a native client that already holds a token from the customer's identity
provider trades it for the proxy-API credential without a browser round trip."""
_IssuedTokenType = Literal["urn:ietf:params:oauth:token-type:access_token"]
ACCESS_TOKEN_TOKEN_TYPE: Final[_IssuedTokenType] = "urn:ietf:params:oauth:token-type:access_token"
SUBJECT_TOKEN_TYPES: Final = frozenset(
{
"urn:ietf:params:oauth:token-type:jwt",
"urn:ietf:params:oauth:token-type:id_token",
ACCESS_TOKEN_TOKEN_TYPE,
}
)
class SubjectIdentity(BaseModel):
model_config = ConfigDict(frozen=True)
user_id: str = Field(min_length=1)
team_id: str | None = None
class SubjectTokenRefusal(BaseModel):
model_config = ConfigDict(frozen=True)
error: Literal["unsupported_grant_type", "invalid_request", "temporarily_unavailable"]
description: str = Field(min_length=1)
class ExchangeSubjectToken(Protocol):
"""Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the
IdP token the way the proxy's own JWT auth does and names the litellm user and team it
stands for, or says why this gateway will not take it."""
def __call__(self, subject_token: str, request: Request, /) -> Awaitable[SubjectIdentity | SubjectTokenRefusal]: ...
class ConsentTeam(BaseModel):
model_config = ConfigDict(frozen=True)
team_id: str = Field(min_length=1)
@ -213,6 +259,12 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr
return "unresolvable"
async def _refuse_subject_token(subject_token: str, request: Request) -> SubjectTokenRefusal:
return SubjectTokenRefusal(
error="unsupported_grant_type", description="this gateway is not configured to exchange IdP tokens"
)
async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState:
return "unavailable"
@ -318,7 +370,9 @@ def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None:
return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY)
async def register_aggregate_client(request: Request, request_body: Mapping[str, object]) -> Response:
async def register_aggregate_client(
request: Request, request_body: Mapping[str, object], token_exchange_available: bool
) -> Response:
"""RFC 7591 dynamic registration against the gateway itself, statelessly.
Only ``redirect_uris`` is authoritative; every client is registered as a public
@ -382,7 +436,7 @@ async def register_aggregate_client(request: Request, request_body: Mapping[str,
"client_id_issued_at": int(now.timestamp()),
"redirect_uris": list(raw_uris),
"token_endpoint_auth_method": "none",
"grant_types": ["authorization_code", "refresh_token"],
"grant_types": list(supported_grant_types(token_exchange_available)),
"response_types": ["code"],
},
)
@ -580,7 +634,7 @@ class NativeClientAuthContract(TypedDict):
revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]]
def native_client_auth_contract(request: Request) -> NativeClientAuthContract:
def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract:
"""The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a
native client (in any language) needs to run the sign-in without reading LiteLLM
source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter
@ -595,7 +649,7 @@ def native_client_auth_contract(request: Request) -> NativeClientAuthContract:
"revocation_endpoint": f"{base_url}/revoke",
"resource": base_url,
"response_types_supported": ("code",),
"grant_types_supported": ("authorization_code", "refresh_token"),
"grant_types_supported": supported_grant_types(token_exchange_available),
"code_challenge_methods_supported": ("S256",),
"token_endpoint_auth_methods_supported": ("none",),
"revocation_endpoint_auth_methods_supported": ("none",),
@ -1033,20 +1087,26 @@ class _ProxyCredentialTokenResponse(TypedDict):
refresh_token: ReadOnly[str]
user_id: ReadOnly[str]
team_id: ReadOnly[str | None]
issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]]
def _proxy_credential_response(
minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime
minted: MintedProxyCredential,
principal: SessionPrincipal,
keys: SessionSigningKeys,
now: datetime,
issued_token_type: _IssuedTokenType | None = None,
) -> Response:
"""The proxy-API token response: the access token is the very credential ``lite
login`` stores (accepted on every proxy route with user and team attribution), and
the refresh token is a gateway-sealed rotating token bound to the team the credential
was minted for, so a renewal keeps the team the user consented to."""
was minted for, so a renewal keeps the team the user consented to. A token exchange
also states ``issued_token_type``, which RFC 8693 section 2.2.1 requires."""
bound_principal: Final = principal.model_copy(update=MappingProxyType({"team_id": minted.team_id}))
refresh: Final = mint_session_refresh_token(bound_principal, keys, now)
if not isinstance(refresh, MintedSessionToken):
return _oauth_error(500, "server_error", "failed to mint the session credential")
body: Final[_ProxyCredentialTokenResponse] = {
credential: Final[_ProxyCredentialTokenResponse] = {
"access_token": minted.key,
"token_type": "Bearer",
"expires_in": minted.expires_in,
@ -1054,7 +1114,10 @@ def _proxy_credential_response(
"user_id": minted.user_id,
"team_id": minted.team_id,
}
return JSONResponse(status_code=200, content=body, headers=TOKEN_NO_CACHE_HEADERS)
if issued_token_type is None:
return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS)
exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type}
return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS)
def _reload_failure_response(failure: ReloadUserFailure) -> Response:
@ -1073,6 +1136,16 @@ def _reload_failure_response(failure: ReloadUserFailure) -> Response:
assert_never(failure)
def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response:
match refusal.error:
case "temporarily_unavailable":
return _oauth_error(503, refusal.error, refusal.description)
case "unsupported_grant_type" | "invalid_request":
return _oauth_error(400, refusal.error, refusal.description)
case _:
assert_never(refusal.error)
def _mint_failure_response(failure: ProxyCredentialMintFailure) -> Response:
match failure:
case "not_a_member":
@ -1116,11 +1189,16 @@ async def aggregate_token(
cache: DualCache,
resource: str | None = None,
mint_proxy_credential: MintProxyCredential = _refuse_proxy_credential,
subject_token: str | None = None,
subject_token_type: str | None = None,
requested_token_type: str | None = None,
exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token,
) -> Response:
"""The aggregate token verb: authorization_code and refresh_token grants for the
identity-only session pair, or for the proxy-API credential when the grant was issued
with that audience. Every path re-validates the litellm user live before minting, so a
deactivated user cannot obtain or renew a session."""
with that audience, and the RFC 8693 token exchange that turns an IdP token straight
into the proxy-API credential. Every path re-validates the litellm user live before
minting, so a deactivated user cannot obtain or renew a session."""
if master_key is None:
verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured")
return _oauth_error(500, "server_error", "the gateway has no master key configured")
@ -1159,7 +1237,20 @@ async def aggregate_token(
now=now,
issue=issue,
)
return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token")
if grant_type == TOKEN_EXCHANGE_GRANT_TYPE:
return await _token_exchange_grant(
subject_token=subject_token,
subject_token_type=subject_token_type,
requested_token_type=requested_token_type,
client_id=client_id,
exchange_subject_token=exchange_subject_token,
issue=issue,
)
return _oauth_error(
400,
"unsupported_grant_type",
f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}",
)
class _GrantIssuer:
@ -1211,10 +1302,9 @@ class _GrantIssuer:
async def _issue_proxy_credential(
self, principal: SessionPrincipal, claim_key: str, claim_ttl_seconds: int, replayed: str
) -> Response:
if self._resource is not None and not is_proxy_api_resource(self._request, self._resource):
return _oauth_error(
400, "invalid_target", "resource does not match the proxy API this grant was issued for"
)
target_refusal: Final = self._proxy_api_target_refusal()
if target_refusal is not None:
return target_refusal
minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id)
if not isinstance(minted, MintedProxyCredential):
return _mint_failure_response(minted)
@ -1223,6 +1313,33 @@ class _GrantIssuer:
return refusal
return _proxy_credential_response(minted, principal, self._keys, self._now)
async def exchange(
self, subject_token: str, client_id: str, exchange_subject_token: ExchangeSubjectToken
) -> Response:
"""The RFC 8693 tail: prove the IdP token, then mint. No single-use marker, because
the subject token stays a valid proof for as long as the IdP says it is and every
exchange mints a fresh credential and refresh token of its own."""
target_refusal: Final = self._proxy_api_target_refusal()
if target_refusal is not None:
return target_refusal
identity: Final = await exchange_subject_token(subject_token, self._request)
if isinstance(identity, SubjectTokenRefusal):
return _subject_token_refusal_response(identity)
principal: Final = SessionPrincipal(
user_id=identity.user_id, client_id=client_id, audience=PROXY_API_AUDIENCE, team_id=identity.team_id
)
minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id)
if not isinstance(minted, MintedProxyCredential):
return _mint_failure_response(minted)
return _proxy_credential_response(
minted, principal, self._keys, self._now, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE
)
def _proxy_api_target_refusal(self) -> Response | None:
if self._resource is None or is_proxy_api_resource(self._request, self._resource):
return None
return _oauth_error(400, "invalid_target", "resource does not match the proxy API this grant was issued for")
async def _claim_refusal(self, claim_key: str, claim_ttl_seconds: int, replayed: str) -> Response | None:
return _claim_refusal(
await self._guard.claim(claim_key, claim_ttl_seconds), replayed=_oauth_error(400, "invalid_grant", replayed)
@ -1297,6 +1414,32 @@ async def _refresh_token_grant(
)
async def _token_exchange_grant(
subject_token: str | None,
subject_token_type: str | None,
requested_token_type: str | None,
client_id: str,
exchange_subject_token: ExchangeSubjectToken,
issue: _GrantIssuer,
) -> Response:
"""RFC 8693 token exchange for a registered native client that already holds an IdP
token: the gateway proves the token the way its JWT auth does and answers with the
proxy-API credential, so a fresh laptop with only an IdP login gets a gateway key
without a browser round trip. The client must be registered because the refresh token
in the answer is bound to it."""
if not is_gateway_dcr_client_id(client_id) or open_gateway_dcr_client(client_id) is None:
return _oauth_error(401, "invalid_client", "unknown or malformed client_id")
if not subject_token or not subject_token_type:
return _oauth_error(400, "invalid_request", "subject_token and subject_token_type are required")
if subject_token_type not in SUBJECT_TOKEN_TYPES:
return _oauth_error(
400, "invalid_request", f"subject_token_type must be one of {', '.join(sorted(SUBJECT_TOKEN_TYPES))}"
)
if requested_token_type is not None and requested_token_type != ACCESS_TOKEN_TOKEN_TYPE:
return _oauth_error(400, "invalid_request", f"requested_token_type must be {ACCESS_TOKEN_TOKEN_TYPE}")
return await issue.exchange(subject_token, client_id, exchange_subject_token)
async def revoke_refresh_token(token: str, client_id: str, master_key: str | None, cache: DualCache) -> Response:
"""RFC 7009 revocation for the gateway's refresh tokens: burn the presented token's
``jti`` so neither the holder nor a thief can rotate it again. Access tokens are

217
litellm/proxy/_experimental/mcp_server/idp_token_exchange.py Normal file
View file

@ -0,0 +1,217 @@
"""The identity-provider side of the RFC 8693 token exchange on ``POST /token``: a native
client that already holds a JWT from the customer's IdP trades it for the same proxy-API
credential ``lite login`` stores, proven by the proxy's own JWT auth (signature, claims,
and the user and team sync it performs), so no browser round trip is needed."""
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from typing import Final, Literal, Protocol
from fastapi import HTTPException, Request
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal
from litellm.proxy._types import JWTAuthBuilderResult, ProxyException
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
EXCHANGE_ROUTE: Final = "/token"
REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth"
SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = (
"the gateway could not verify subject_token because its identity provider or database is unavailable; retry"
)
SUBJECT_TOKEN_CHECK_FAULTED: Final = (
"the gateway could not verify subject_token because its database reported a fault that is not a transient "
"outage; retrying will not help until the gateway deployment is repaired"
)
GatewayOutage = Literal["retryable", "faulted"]
@dataclass(frozen=True, slots=True)
class TokenExchangePrerequisites:
"""The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT
bearer, plus the JWT-to-virtual-key mapping it consults first: a gateway that maps
tokens authenticates a JWT as its mapped key, with that key's models and budget, or
refuses an unmapped one, and the exchange proves the token through ``auth_builder``
alone, so it would mint the user's own credential past that policy. Discovery and
registration advertise the exchange grant only when every gate holds, and an exchange
attempt is refused naming the first one that does not."""
jwt_auth_enabled: bool
has_database: bool
licensed: bool
maps_jwts_to_virtual_keys: bool
@property
def available(self) -> bool:
return self.jwt_auth_enabled and self.has_database and self.licensed and not self.maps_jwts_to_virtual_keys
def refusal(self) -> SubjectTokenRefusal | None:
if not self.jwt_auth_enabled:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens",
)
if not self.has_database:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="this gateway has no database, so it cannot exchange IdP tokens",
)
if not self.licensed:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="JWT auth is an enterprise only feature; no license is set",
)
if self.maps_jwts_to_virtual_keys:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="this gateway maps IdP tokens to virtual keys, which the exchange does not serve",
)
return None
def read_token_exchange_prerequisites() -> TokenExchangePrerequisites:
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
general_settings,
jwt_handler,
premium_user,
prisma_client,
)
return TokenExchangePrerequisites(
jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True,
has_database=prisma_client is not None,
licensed=premium_user is True,
maps_jwts_to_virtual_keys=_maps_jwts_to_virtual_keys(jwt_handler),
)
def _maps_jwts_to_virtual_keys(jwt_handler: JWTHandler) -> bool:
if not hasattr(jwt_handler, "litellm_jwtauth"):
return False
return jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured()
def token_exchange_available() -> bool:
return read_token_exchange_prerequisites().available
class AuthorizeSubjectToken(Protocol):
"""Injected JWT authorization ``(subject_token, request_headers)``: the proxy's
``JWTAuthManager.auth_builder`` in production, which raises when the token is not
acceptable and otherwise names the user and team it resolved."""
def __call__(
self, subject_token: str, request_headers: Mapping[str, str], /
) -> Awaitable[JWTAuthBuilderResult]: ...
async def exchange_idp_subject_token(subject_token: str, request: Request) -> SubjectIdentity | SubjectTokenRefusal:
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
general_settings,
jwt_handler,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
async def authorize(token: str, request_headers: Mapping[str, str]) -> JWTAuthBuilderResult:
return await JWTAuthManager.auth_builder(
api_key=token,
jwt_handler=jwt_handler,
request_data={},
general_settings=general_settings,
route=EXCHANGE_ROUTE,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
request_headers=request_headers,
request_method="POST",
)
return await identity_from_subject_token(
subject_token,
request_headers=request.headers,
prerequisites=read_token_exchange_prerequisites(),
is_jwt=jwt_handler.is_jwt,
authorize=authorize,
)
async def identity_from_subject_token(
subject_token: str,
request_headers: Mapping[str, str],
prerequisites: TokenExchangePrerequisites,
is_jwt: Callable[[str], bool],
authorize: AuthorizeSubjectToken,
) -> SubjectIdentity | SubjectTokenRefusal:
"""Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the
proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which
RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token, and a
check the gateway could not complete (the IdP's JWKS unreachable with no cached copy,
the auth database down) as ``temporarily_unavailable``, so the client retries instead
of treating a valid token as bad, worded by whether retrying can help. The reason stays
in the proxy log: this endpoint is public and JWT auth's own wording can name the JWKS
URL it fetched or quote the IdP's response."""
unmet: Final = prerequisites.refusal()
if unmet is not None:
return unmet
if not is_jwt(subject_token):
return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT")
try:
result: Final = await authorize(subject_token, request_headers)
except HTTPException as denied:
return _refusal_for(denied, denied.detail)
except ProxyException as denied:
return _refusal_for(denied, denied.message)
except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures
return _refusal_for(denied, denied)
user_id: Final = result["user_id"]
if user_id is None:
return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows")
return SubjectIdentity(user_id=user_id, team_id=result["team_id"])
def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal:
outage: Final = _gateway_could_not_verify(denied)
if outage is None:
verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason)
return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
verbose_proxy_logger.error("token exchange could not verify a subject_token, %s: %s", outage, reason)
return SubjectTokenRefusal(error="temporarily_unavailable", description=_check_unavailable_description(outage))
def _check_unavailable_description(outage: GatewayOutage) -> str:
match outage:
case "retryable":
return SUBJECT_TOKEN_CHECK_UNAVAILABLE
case "faulted":
return SUBJECT_TOKEN_CHECK_FAULTED
case _:
assert_never(outage)
def _gateway_could_not_verify(denied: Exception) -> GatewayOutage | None:
"""A database fault anywhere in the chain (``get_user_object`` wraps prisma failures in a
bare ``ValueError``) or a 5xx from JWT auth (the IdP's JWKS unreachable with no cached
copy) is the gateway failing, not the token. A fault retrying cannot clear (a missing or
version-skewed query engine) is named as such, the way the mint path words it, so the
client is not told to wait on a deployment that needs repair."""
fault: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(denied)
if fault is not None:
return "faulted" if PrismaDBExceptionHandler.is_permanent_database_fault(fault) else "retryable"
return "retryable" if _is_server_error(denied) else None
def _is_server_error(denied: Exception) -> bool:
match denied:
case HTTPException(status_code=status_code):
return status_code >= 500
case ProxyException(code=code):
return code.isdigit() and int(code) >= 500
case _:
return False

16
litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py
View file

@ -16,7 +16,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
ReloadUserFailure,
)
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, effective_user_role
from litellm.proxy.management_endpoints.ui_sso import (
CliSsoTeamDetail,
fetch_cli_sso_team_details,
@ -42,7 +42,7 @@ async def mint_proxy_credential(
user_id: str, team_id: str | None
) -> MintedProxyCredential | ProxyCredentialMintFailure:
"""Mint the ``lite login`` credential for a consented grant. Membership is checked
live, so a team the user left between consent and redemption (or between refreshes)
live against the database row, so a team the user left between consent and redemption (or between refreshes)
refuses the grant instead of minting a credential attributed to a team they are no
longer on. The team is exactly the one the consent page sealed into the grant; nothing
is picked on the user's behalf here, so a refresh can never move the credential, and a
@ -51,12 +51,12 @@ async def mint_proxy_credential(
posting the consent form without one. Memberships whose team rows are gone count as no
team at all, the way ``lite login`` treats them, so they can never lock a user out. The
user row handed to the minter carries no team list, exactly like ``lite login``'s, so
the minter's own first-team fallback stays inert."""
user: Final = await load_active_user_by_id(user_id)
the minter's own first-team fallback stays inert. The credential carries the role the
proxy already enforces for the user on every request, so a row with no role (JWT auth's
upsert writes none) mints as an internal user instead of being refused."""
user: Final = await load_active_user_by_id(user_id, source="database")
if isinstance(user, str):
return user
if user.user_role is None:
return "no_active_key"
if team_id is not None and team_id not in user.teams:
return "not_a_member"
details: Final = await _team_details(user.teams) if user.teams else ()
@ -68,7 +68,9 @@ async def mint_proxy_credential(
if selected is None:
return "not_a_member"
key: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info=LiteLLM_UserTable(user_id=user.user_id, user_role=user.user_role, models=user.models),
user_info=LiteLLM_UserTable(
user_id=user.user_id, user_role=effective_user_role(user.user_role).value, models=user.models
),
team_id=team_id,
team_alias=selected.team_alias,
team_models=selected.team_models,

3
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -214,6 +214,7 @@ if MCP_AVAILABLE:
filter_tools_by_allowed_tools,
filter_tools_by_key_team_permissions,
fire_mcp_tool_call_failure_logging,
reject_disallowed_mcp_client,
)
########################################################
@ -896,6 +897,7 @@ if MCP_AVAILABLE:
MCPRequestHandler,
)
reject_disallowed_mcp_client(request.headers, user_api_key_dict)
try:
mcp_server_name = _as_query_str(mcp_server_name)
toolset_name = _as_query_str(toolset_name)
@ -1099,6 +1101,7 @@ if MCP_AVAILABLE:
proxy_logging_obj,
)
reject_disallowed_mcp_client(request.headers, user_api_key_dict)
try:
user_api_key_dict = await acting_user_auth(user_api_key_dict)
data = await request.json()

27
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -47,6 +47,11 @@ from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
cache_byok_credential,
get_cached_byok_credential,
)
from litellm.proxy._experimental.mcp_server.client_allowlist import (
MCPClientAllowlist,
check_mcp_client_allowed,
load_mcp_client_allowlist,
)
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
get_request_base_url,
)
@ -74,6 +79,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
get_route_relative_request_path,
well_known_root_suffix,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import is_ui_session_credential
from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_DESCRIPTION,
LITELLM_MCP_SERVER_NAME,
@ -3673,6 +3679,25 @@ if MCP_AVAILABLE:
mcp_servers_from_path = [servers_and_path]
return mcp_servers_from_path
def _load_mcp_client_allowlist() -> MCPClientAllowlist | None:
from litellm.proxy.proxy_server import general_settings
return load_mcp_client_allowlist(general_settings)
def reject_disallowed_mcp_client(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth | None) -> None:
"""Gate every MCP tool surface on ``mcp_allowed_clients``; the dashboard's own session is not a client app."""
if user_api_key_auth is not None and is_ui_session_credential(user_api_key_auth):
return
rejection: Final = check_mcp_client_allowed(
allowlist=_load_mcp_client_allowlist(),
jwt_claims=user_api_key_auth.jwt_claims if user_api_key_auth is not None else None,
headers=headers,
)
if rejection is None:
return
verbose_logger.warning("Rejected MCP request from a disallowed client application: %s", rejection.details)
raise HTTPException(status_code=403, detail=rejection.response_body)
async def extract_mcp_auth_context(scope, path):
"""
Extracts mcp_servers from the path and processes the MCP request for auth context.
@ -4524,6 +4549,7 @@ if MCP_AVAILABLE:
oauth2_headers,
raw_headers,
) = await extract_mcp_auth_context(scope, path)
reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth)
scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1
# Extract client IP for MCP access control
@ -4852,6 +4878,7 @@ if MCP_AVAILABLE:
oauth2_headers,
raw_headers,
) = await extract_mcp_auth_context(scope, path)
reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth)
scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1
# Extract client IP for MCP access control

5
litellm/proxy/_lazy_features.py
View file

@ -233,6 +233,11 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
module_path="litellm.proxy.anthropic_endpoints.skills_endpoints",
path_prefixes=("/v1/skills", "/skills"),
),
LazyFeature(
name="claude_code_gateway",
module_path="litellm.proxy.anthropic_endpoints.gateway_endpoints",
path_prefixes=("/claude_code_gateway",),
),
LazyFeature(
name="langfuse_passthrough",
module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints",

72
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -5235,6 +5235,12 @@
}
}
},
"claude_code_gateway": {
"components": {
"schemas": {}
},
"paths": {}
},
"claude_code_marketplace": {
"components": {
"schemas": {
@ -23678,6 +23684,17 @@
],
"title": "Refresh Token"
},
"requested_token_type": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Requested Token Type"
},
"resource": {
"anyOf": [
{
@ -23699,6 +23716,28 @@
}
],
"title": "Scope"
},
"subject_token": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Subject Token"
},
"subject_token_type": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Subject Token Type"
}
},
"required": [
@ -23752,6 +23791,17 @@
],
"title": "Refresh Token"
},
"requested_token_type": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Requested Token Type"
},
"resource": {
"anyOf": [
{
@ -23773,6 +23823,28 @@
}
],
"title": "Scope"
},
"subject_token": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Subject Token"
},
"subject_token_type": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Subject Token Type"
}
},
"required": [

38
litellm/proxy/_types.py
View file

@ -34,6 +34,7 @@ from litellm.types.llms.openai import (
ResponsesAPIResponse,
)
from litellm.types.mcp import (
MCPAllowedClient,
MCPAuth,
MCPAuthType,
MCPCredentials,
@ -511,6 +512,8 @@ class LiteLLMRoutes(enum.Enum):
anthropic_routes = [
"/v1/messages",
"/v1/messages/count_tokens",
"/claude_code_gateway/v1/messages",
"/claude_code_gateway/v1/messages/count_tokens",
"/v1/skills",
"/v1/skills/{skill_id}",
"/claude-code/marketplace.json",
@ -531,6 +534,7 @@ class LiteLLMRoutes(enum.Enum):
"/mcp-rest/tools/call",
"/v1/mcp/tools",
"/introspect",
"/token",
]
# MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS.
@ -888,6 +892,11 @@ class LiteLLMRoutes(enum.Enum):
# of; a caller who administers none gets an empty result set.
"/organization/daily/activity",
"/user/available_roles", # read-only role metadata; any authenticated user may read
# Claude Code gateway: the signed-in CLI fetches its managed settings and posts its own telemetry
"/claude_code_gateway/managed/settings",
"/claude_code_gateway/v1/metrics",
"/claude_code_gateway/v1/logs",
"/claude_code_gateway/v1/traces",
"/user/list", # org admins checked in endpoint; non-admins get 403
"/management/v1/users/bulk_delete", # proxy admins delete anyone, org admins only their orgs' users; others 403
"/model/{model_id}/update",
@ -897,6 +906,9 @@ class LiteLLMRoutes(enum.Enum):
# Project read routes - endpoint scopes results to caller's teams (non-admin)
"/project/list",
"/project/info",
# Project write routes - endpoint checks team admin + team_admin_editable_team_fields "projects"
"/project/new",
"/project/update",
# Endpoint enforces proxy-admin vs team-admin model access itself.
"/health/test_connection",
# Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges
@ -2600,6 +2612,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="opt-in to RFC 8628 verification_uri_complete for the CLI SSO device flow, pre-filling the user_code in the browser. Off by default; intended for same-host clients where the device that starts the flow and the browser run on the same machine",
)
enable_claude_code_gateway: bool | None = Field(
None,
description="serve the Claude Code gateway protocol (https://code.claude.com/docs/en/claude-apps-gateway) under /claude_code_gateway: OAuth device-flow sign-in reusing proxy SSO, plus managed settings and OTLP telemetry ingestion. Off by default",
)
claude_code_gateway_managed_settings: dict[str, Any] | None = Field(
None,
description="Claude Code managed-settings.json served verbatim at the gateway's /claude_code_gateway/managed/settings endpoint. When unset the endpoint returns 404 (no managed policy)",
)
database_url: str | None = Field(
None,
description="connect to a postgres db - needed for generating temporary keys + tracking spend / key",
@ -2902,6 +2922,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).",
)
mcp_allowed_clients: list[MCPAllowedClient] | None = Field(
None,
description="MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted.",
)
mcp_client_id_header: str | None = Field(
None,
description="Request header whose value names the calling MCP client application (for example 'x-mcp-client') for callers that did not authenticate with a JWT, used only while mcp_allowed_clients is set. The client picks this value itself, so it is a policy control rather than a security boundary; prefer litellm_jwtauth.mcp_client_id_jwt_field where callers use JWTs.",
)
mcp_trusted_proxy_ranges: list[str] | None = Field(
None,
description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For and X-Forwarded-* origin headers are only trusted from these IPs.",
@ -5118,6 +5146,15 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
"then agent_name, and the request is rejected when it matches neither."
),
)
mcp_client_id_jwt_field: str | None = Field(
default=None,
description=(
"The field in the JWT token that identifies the MCP client application (harness) making the request, "
"e.g. 'azp' or 'client_id'. Supports dot notation. Only consulted while general_settings.mcp_allowed_clients "
"is set: the claim value must be listed there or the MCP request is rejected with 403. Distinct from "
"agent_id_jwt_field, which identifies an AI agent rather than the client software."
),
)
public_key_ttl: float = 600
public_key_stale_ttl: float = Field(
default=DEFAULT_JWKS_STALE_TTL,
@ -5403,6 +5440,7 @@ class DBSpendUpdateTransactions(TypedDict):
team_member_list_transactions: dict[str, float] | None
org_list_transactions: dict[str, float] | None
org_member_list_transactions: ReadOnly[dict[str, float] | None]
project_list_transactions: ReadOnly[dict[str, float] | None]
tag_list_transactions: dict[str, float] | None
agent_list_transactions: dict[str, float] | None
model_access_group_list_transactions: ReadOnly[dict[str, float] | None]

397
litellm/proxy/anthropic_endpoints/gateway_endpoints.py Normal file
View file

@ -0,0 +1,397 @@
"""
Claude Code gateway protocol.
Implements the wire contract the Claude Code CLI uses to talk to a gateway:
OAuth 2.0 device-authorization sign-in (RFC 8414 / RFC 8628), inference via the
Anthropic Messages API, managed settings, and OTLP telemetry ingestion. See
https://code.claude.com/docs/en/claude-apps-gateway.
Everything lives under the ``/claude_code_gateway`` base so operators point
Claude Code at ``https://<proxy-host>/claude_code_gateway`` via ``/login``. The
device flow reuses the proxy's existing SSO login machinery: the browser leg is
served by ``/sso/key/generate`` and the shared ``cli_sso_session_cache`` flow,
so the bearer token minted here is the same session JWT the LiteLLM CLI uses and
is accepted by every bearer-authenticated proxy route.
"""
import hashlib
import json
import secrets
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import (
CLI_JWT_EXPIRATION_HOURS,
CLI_SSO_SESSION_TTL_SECONDS,
LITELLM_CLI_SOURCE_IDENTIFIER,
)
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles
from litellm.proxy.anthropic_endpoints.endpoints import anthropic_response, count_tokens
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.http_parsing_utils import _safe_set_request_parsed_body
from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail
GATEWAY_PREFIX: Final = "/claude_code_gateway"
_DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code"
_REFRESH_TOKEN_GRANT: Final = "refresh_token"
_DEVICE_CODE_SEPARATOR: Final = "."
_DEVICE_POLL_INTERVAL_SECONDS: Final = 5
_SECONDS_PER_HOUR: Final = 3600
_MANAGED_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object])
_NO_SETTINGS: Final = MappingProxyType({})
_POST_ONLY: Final = ["POST"] # mutable-ok: FastAPI's add_api_route only accepts a list of methods
class _GatewaySessionData(BaseModel):
user_id: str
user_role: LitellmUserRoles
models: list[str] = Field(default_factory=list)
teams: tuple[str, ...] = ()
team_details: object | None = None
@dataclass(frozen=True, slots=True)
class _GatewayLogin:
user_info: LiteLLM_UserTable
team_id: str | None
team: CliSsoTeamDetail
class _OAuthErrorBody(BaseModel):
error: str
error_description: str | None = None
class _AuthorizationServerMetadata(BaseModel):
issuer: str
device_authorization_endpoint: str
token_endpoint: str
grant_types_supported: tuple[str, ...]
class _DeviceAuthorizationBody(BaseModel):
device_code: str
user_code: str
verification_uri: str
verification_uri_complete: str | None = None
expires_in: int
interval: int
class _AccessTokenBody(BaseModel):
access_token: str
expires_in: int
token_type: str = "Bearer"
class _ManagedSettingsBody(BaseModel):
uuid: str
checksum: str
settings: dict[str, object]
def _general_settings() -> Mapping[str, object]:
from litellm.proxy.proxy_server import general_settings
return general_settings or _NO_SETTINGS
def _is_gateway_enabled() -> bool:
return bool(_general_settings().get("enable_claude_code_gateway", False))
def ensure_gateway_enabled() -> None:
from fastapi import HTTPException
if not _is_gateway_enabled():
raise HTTPException(status_code=404, detail="Claude Code gateway is not enabled")
def _managed_settings() -> dict[str, object] | None:
settings: Final[object] = _general_settings().get("claude_code_gateway_managed_settings")
if not isinstance(settings, dict):
return None
return _MANAGED_SETTINGS_ADAPTER.validate_python(settings)
@dataclass(frozen=True, slots=True)
class _OAuthError:
status_code: int
error: str
description: str | None = None
def _oauth_error_response(err: _OAuthError) -> JSONResponse:
body: Final = _OAuthErrorBody(error=err.error, error_description=err.description)
return JSONResponse(status_code=err.status_code, content=body.model_dump(exclude_none=True))
router: Final = APIRouter(
prefix=GATEWAY_PREFIX,
tags=["Claude Code gateway"], # mutable-ok: FastAPI's APIRouter only accepts a list of tags
)
_GATEWAY_ENABLED: Final = (Depends(ensure_gateway_enabled),)
_AUTHENTICATED: Final = (Depends(user_api_key_auth),)
router.add_api_route(
"/v1/messages",
anthropic_response,
methods=_POST_ONLY,
dependencies=_GATEWAY_ENABLED,
include_in_schema=False,
)
router.add_api_route(
"/v1/messages/count_tokens",
count_tokens,
methods=_POST_ONLY,
dependencies=_GATEWAY_ENABLED,
include_in_schema=False,
)
@router.get("/.well-known/oauth-authorization-server", include_in_schema=False)
async def oauth_authorization_server(request: Request) -> JSONResponse:
if not _is_gateway_enabled():
return _oauth_error_response(_OAuthError(status_code=404, error="not_found"))
from litellm.proxy.utils import get_custom_url
request_base_url: Final = str(request.base_url)
metadata: Final = _AuthorizationServerMetadata(
issuer=get_custom_url(request_base_url=request_base_url, route="claude_code_gateway"),
device_authorization_endpoint=get_custom_url(
request_base_url=request_base_url, route="claude_code_gateway/oauth/device_authorization"
),
token_endpoint=get_custom_url(request_base_url=request_base_url, route="claude_code_gateway/oauth/token"),
grant_types_supported=(_DEVICE_CODE_GRANT, _REFRESH_TOKEN_GRANT),
)
return JSONResponse(content=metadata.model_dump())
@router.post("/oauth/device_authorization", include_in_schema=False)
async def device_authorization(request: Request) -> JSONResponse:
from urllib.parse import urlencode
from litellm.proxy.management_endpoints.ui_sso import (
_check_cli_sso_start_rate_limit, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
_cli_sso_verification_uri_complete_enabled, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
_generate_cli_sso_user_code, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
_hash_cli_sso_secret, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
_normalize_cli_sso_user_code, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
_set_cli_sso_flow, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
)
from litellm.proxy.proxy_server import cli_sso_session_cache
from litellm.proxy.utils import get_custom_url
if not _is_gateway_enabled():
return _oauth_error_response(_OAuthError(status_code=404, error="not_found"))
_check_cli_sso_start_rate_limit(
request=request,
cache=cli_sso_session_cache,
use_x_forwarded_for=bool(_general_settings().get("use_x_forwarded_for", False)),
)
login_id: Final = f"cli-{secrets.token_urlsafe(24)}"
poll_secret: Final = secrets.token_urlsafe(32)
user_code: Final = _generate_cli_sso_user_code()
flow: Final = { # mutable-ok: the shared CLI SSO cache entry is a dict the browser leg mutates
"poll_secret_hash": _hash_cli_sso_secret(poll_secret),
"user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)),
"sso_complete": False,
"user_code_verified": False,
"session_data": None,
}
_set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow)
request_base_url: Final = str(request.base_url)
verification_uri: Final = get_custom_url(request_base_url=request_base_url, route="sso/key/generate")
query: Final = MappingProxyType({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": login_id})
body: Final = _DeviceAuthorizationBody(
device_code=f"{login_id}{_DEVICE_CODE_SEPARATOR}{poll_secret}",
user_code=user_code,
verification_uri=f"{verification_uri}?{urlencode(query)}",
verification_uri_complete=(
f"{verification_uri}?{urlencode(MappingProxyType({**query, 'user_code': user_code}))}"
if _cli_sso_verification_uri_complete_enabled()
else None
),
expires_in=CLI_SSO_SESSION_TTL_SECONDS,
interval=_DEVICE_POLL_INTERVAL_SECONDS,
)
return JSONResponse(content=body.model_dump(exclude_none=True))
def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError:
from litellm.proxy.management_endpoints.ui_sso import selected_cli_sso_team_detail
try:
session_data: Final = _GatewaySessionData.model_validate(flow.get("session_data"))
except ValidationError as err:
verbose_proxy_logger.warning("Claude Code gateway login session is malformed: %s", err)
return _OAuthError(
status_code=400, error="invalid_grant", description="The login session is malformed; sign in again"
)
team_id: Final = session_data.teams[0] if session_data.teams else None
selected_team: Final = selected_cli_sso_team_detail(team_details=session_data.team_details, team_id=team_id)
if selected_team is None:
return _OAuthError(
status_code=400,
error="invalid_grant",
description=f"Could not resolve the model grants for team {team_id}; sign in again",
)
user_info: Final = LiteLLM_UserTable(
user_id=session_data.user_id,
user_role=session_data.user_role.value,
models=session_data.models,
)
return _GatewayLogin(user_info=user_info, team_id=team_id, team=selected_team)
def _mint_access_token(login: _GatewayLogin) -> str:
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
return ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info=login.user_info,
team_id=login.team_id,
team_alias=login.team.team_alias,
team_models=login.team.team_models,
team_model_aliases=login.team.team_model_aliases,
max_budget=None,
)
async def _claim_device_code(login_id: str, cache: DualCache) -> bool:
from litellm.proxy.management_endpoints.ui_sso import (
_get_cli_sso_flow_cache_key, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
)
claims: Final = await cache.async_increment_cache(
key=f"{_get_cli_sso_flow_cache_key(login_id)}:claimed",
value=1,
ttl=CLI_SSO_SESSION_TTL_SECONDS,
)
return claims == 1
async def _handle_device_code_grant(device_code: str | None) -> JSONResponse:
from fastapi import HTTPException
from litellm.proxy.management_endpoints.ui_sso import (
_get_cli_sso_flow_cache_key, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
_get_cli_sso_flow_or_raise, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
_verify_cli_sso_poll_secret, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
)
from litellm.proxy.proxy_server import cli_sso_session_cache
if not device_code:
return _oauth_error_response(
_OAuthError(status_code=400, error="invalid_request", description="device_code is required")
)
login_id, _, poll_secret = device_code.partition(_DEVICE_CODE_SEPARATOR)
try:
flow: Final = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cli_sso_session_cache)
except HTTPException:
return _oauth_error_response(_OAuthError(status_code=400, error="expired_token"))
if not _verify_cli_sso_poll_secret(flow, poll_secret):
return _oauth_error_response(_OAuthError(status_code=400, error="expired_token"))
if not flow.get("sso_complete") or not flow.get("user_code_verified"):
return _oauth_error_response(_OAuthError(status_code=400, error="authorization_pending"))
login: Final = _validate_login(flow)
if isinstance(login, _OAuthError):
return _oauth_error_response(login)
access_token: Final = _mint_access_token(login)
if not await _claim_device_code(login_id, cli_sso_session_cache):
return _oauth_error_response(_OAuthError(status_code=400, error="expired_token"))
await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(login_id))
body: Final = _AccessTokenBody(access_token=access_token, expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR)
return JSONResponse(content=body.model_dump())
@router.post("/oauth/token", include_in_schema=False)
async def oauth_token(request: Request) -> JSONResponse:
if not _is_gateway_enabled():
return _oauth_error_response(_OAuthError(status_code=404, error="not_found"))
form: Final = await request.form()
grant_type: Final = form.get("grant_type")
if grant_type == _DEVICE_CODE_GRANT:
device_code: Final = form.get("device_code")
return await _handle_device_code_grant(device_code if isinstance(device_code, str) else None)
if grant_type == _REFRESH_TOKEN_GRANT:
return _oauth_error_response(
_OAuthError(
status_code=401,
error="invalid_grant",
description="This gateway does not issue refresh tokens; sign in again",
)
)
return _oauth_error_response(
_OAuthError(
status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}"
)
)
@router.get("/managed/settings", include_in_schema=False, dependencies=_AUTHENTICATED)
async def managed_settings(request: Request) -> Response:
ensure_gateway_enabled()
settings: Final = _managed_settings()
if settings is None:
return Response(status_code=404)
canonical: Final = json.dumps(settings, sort_keys=True, separators=(",", ":"))
checksum: Final = "sha256:" + hashlib.sha256(canonical.encode("utf-8")).hexdigest()
etag: Final = f'"{checksum}"'
headers: Final = MappingProxyType({"ETag": etag})
if request.headers.get("If-None-Match") == etag:
return Response(status_code=304, headers=headers)
body: Final = _ManagedSettingsBody(uuid=checksum, checksum=checksum, settings=settings)
return Response(content=body.model_dump_json(), media_type="application/json", headers=headers)
async def _skip_otlp_body_parsing(request: Request) -> None:
_safe_set_request_parsed_body(request=request, parsed_body={})
_OTLP_AUTHENTICATED: Final = (Depends(_skip_otlp_body_parsing), *_AUTHENTICATED)
def _accept_otlp() -> Response:
ensure_gateway_enabled()
return Response(status_code=200)
@router.post("/v1/metrics", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED)
async def otlp_metrics() -> Response:
return _accept_otlp()
@router.post("/v1/logs", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED)
async def otlp_logs() -> Response:
return _accept_otlp()
@router.post("/v1/traces", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED)
async def otlp_traces() -> Response:
return _accept_otlp()

76
litellm/proxy/auth/auth_checks.py
View file

@ -15,10 +15,10 @@ import re
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol
from fastapi import HTTPException, Request, status
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
import litellm
@ -94,6 +94,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
model_access_group_registry_cache_key,
model_access_group_spend_counter_key,
object_permission_cache_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
tag_registry_cache_key,
team_membership_auth_cache_key,
@ -1214,21 +1216,19 @@ async def common_checks(
return True
def effective_user_role(user_role: str | None) -> LitellmUserRoles:
try:
return LitellmUserRoles(user_role)
except ValueError:
return LitellmUserRoles.INTERNAL_USER
def _get_user_role(
user_obj: LiteLLM_UserTable | None,
) -> LitellmUserRoles | None:
if user_obj is None:
return None
_user: Final = user_obj
_user_role: Final = _user.user_role
try:
role: Final = LitellmUserRoles(_user_role)
except ValueError:
return LitellmUserRoles.INTERNAL_USER
return role
return effective_user_role(user_obj.user_role)
def _is_api_route_allowed(
@ -2412,22 +2412,22 @@ def _update_last_db_access_time(key: str, value: object | None, last_db_access_t
last_db_access_time[key] = (value, time.time())
ROLE_BASED_PERMISSIONS_ADAPTER: Final[TypeAdapter[list[RoleBasedPermissions]]] = TypeAdapter(list[RoleBasedPermissions])
def _get_role_based_permissions(
rbac_role: RBAC_ROLES,
general_settings: dict,
general_settings: Mapping[str, object],
key: Literal["models", "routes"],
) -> list[str] | None:
"""
Get the role based permissions from the general settings.
"""
role_based_permissions: Final = cast(
list[RoleBasedPermissions] | None,
general_settings.get("role_permissions", []),
)
if role_based_permissions is None:
configured: Final = general_settings.get("role_permissions")
if configured is None:
return None
for role_based_permission in role_based_permissions:
for role_based_permission in ROLE_BASED_PERMISSIONS_ADAPTER.validate_python(configured):
if role_based_permission.role == rbac_role:
return role_based_permission.models if key == "models" else role_based_permission.routes
@ -2436,7 +2436,7 @@ def _get_role_based_permissions(
def get_role_based_models(
rbac_role: RBAC_ROLES,
general_settings: dict,
general_settings: Mapping[str, object],
) -> list[str] | None:
"""
Get the models allowed for a user role.
@ -2453,7 +2453,7 @@ def get_role_based_models(
def get_role_based_routes(
rbac_role: RBAC_ROLES,
general_settings: dict,
general_settings: Mapping[str, object],
) -> list[str] | None:
"""
Get the routes allowed for a user role.
@ -2575,7 +2575,7 @@ async def get_user_object(
raise Exception("No db connected")
try:
db_access_time_key: Final = f"user_id:{user_id}"
should_check_db: Final = _should_check_db(
should_check_db: Final = bool(check_db_only) or _should_check_db(
key=db_access_time_key,
last_db_access_time=last_db_access_time,
db_cache_expiry=db_cache_expiry,
@ -5680,16 +5680,22 @@ async def _project_max_budget_check(
if project_object.litellm_budget_table is not None:
max_budget = project_object.litellm_budget_table.max_budget
if (
max_budget is not None
and project_object.spend is not None
and math.isfinite(max_budget)
and project_object.spend > max_budget
):
if max_budget is None or max_budget <= 0 or not math.isfinite(max_budget):
return
from litellm.proxy.proxy_server import get_current_spend
project_spend: Final = await get_current_spend(
counter_key=project_spend_counter_key(project_object.project_id),
fallback_spend=project_object.spend or 0.0,
max_budget=max_budget,
)
if project_spend >= max_budget:
if valid_token:
call_info: Final = CallInfo(
token=valid_token.token,
spend=project_object.spend,
spend=project_spend,
max_budget=max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
@ -5705,9 +5711,9 @@ async def _project_max_budget_check(
)
raise litellm.BudgetExceededError(
current_cost=project_object.spend,
current_cost=project_spend,
max_budget=max_budget,
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}",
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}",
entity_type=Litellm_EntityType.PROJECT.value,
entity_id=project_object.project_id,
)
@ -5757,10 +5763,6 @@ async def _project_soft_budget_check(
)
def _project_cache_key(project_id: str) -> str:
return f"project_id:{project_id}"
async def get_project_object(
project_id: str,
prisma_client: PrismaClient | None,
@ -5778,7 +5780,7 @@ async def get_project_object(
return None
# Check cache first
cache_key: Final = _project_cache_key(project_id)
cache_key: Final = project_cache_key(project_id)
deserialized_project: Final = await user_api_key_cache.async_get_cache(
key=cache_key,
model_type=LiteLLM_ProjectTableCachedObj,
@ -5820,7 +5822,7 @@ async def delete_cached_project_object(
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
await evict_and_broadcast(
cache_keys=(_project_cache_key(project_id),),
cache_keys=(project_cache_key(project_id),),
user_api_key_cache=user_api_key_cache,
)

8
litellm/proxy/auth/handle_jwt.py
View file

@ -1867,7 +1867,7 @@ class JWTAuthManager:
@staticmethod
def get_team_id_from_header(
request_headers: dict | None,
request_headers: Mapping[str, str] | None,
allowed_team_ids: set[str],
fallback_to_db_teams: bool = False,
) -> str | None:
@ -2037,7 +2037,7 @@ class JWTAuthManager:
async def _attach_team_from_header_for_admin(
admin_result: JWTAuthBuilderResult,
route: str,
request_headers: dict | None,
request_headers: Mapping[str, str] | None,
jwt_handler: JWTHandler,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
@ -2293,7 +2293,7 @@ class JWTAuthManager:
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
request_headers: dict | None = None,
request_headers: Mapping[str, str] | None = None,
request_method: str | None = None,
) -> JWTAuthBuilderResult:
return await JWTAuthManager.authorize_jwt(
@ -2390,7 +2390,7 @@ class JWTAuthManager:
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
request_headers: dict[str, str] | None = None,
request_headers: Mapping[str, str] | None = None,
request_method: str | None = None,
provisioning: _JWTProvisioning | None = None,
) -> JWTAuthBuilderResult:

24
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -40,6 +40,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
@ -48,6 +50,7 @@ from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import SpendLinkedTable
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
EndUserRepository,
ModelAccessGroupBudgetRepository,
@ -114,6 +117,11 @@ class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol):
def access_group_name(self) -> str: ...
class _ProjectRow(_BudgetLinkedRow, Protocol):
@property
def project_id(self) -> str: ...
class _EndUserRow(_BudgetLinkedRow, Protocol):
@property
def user_id(self) -> str: ...
@ -184,6 +192,14 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]
return (model_access_group_cache_key(row.access_group_name),)
def _project_counter_key(row: _ProjectRow) -> str:
return project_spend_counter_key(row.project_id)
def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]:
return (project_cache_key(row.project_id),)
def _enduser_counter_key(row: _EndUserRow) -> str:
return f"spend:end_user:{row.user_id}"
@ -754,6 +770,11 @@ class ResetBudgetJob:
where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
log_subject="model access groups",
)
projects: Final[tuple[_ProjectRow, ...]] = await self._fetch_linked_rows(
table=ProjectRepository(self.prisma_client).table,
where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
log_subject="projects",
)
rollover_caps: Final[Mapping[str, float]] = MappingProxyType(
{ # mutable-ok: MappingProxyType wraps a one-shot dict comprehension
b.budget_id: cap
@ -786,6 +807,7 @@ class ResetBudgetJob:
(_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps))
for row in model_access_groups
),
*((_project_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in projects),
),
rollover_caps=rollover_caps,
cache_keys=(
@ -794,6 +816,7 @@ class ResetBudgetJob:
*(key for row in orgs for key in _org_cache_keys(row)),
*(key for row in tags for key in _tag_cache_keys(row)),
*(key for row in model_access_groups for key in _model_access_group_cache_keys(row)),
*(key for row in projects for key in _project_cache_keys(row)),
),
)
@ -820,6 +843,7 @@ class ResetBudgetJob:
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE)
_queue_enduser_resets(uow.endusers, cascade)
for budget_id, budget_reset_at in cascade.budget_resets:
uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at)

165
litellm/proxy/common_utils/responses_stream_errors.py Normal file
View file

@ -0,0 +1,165 @@
import time
from collections.abc import Mapping
from http import HTTPStatus
from types import MappingProxyType
from typing import Final
from pydantic import BaseModel, ConfigDict, field_validator
from litellm._logging import redact_internal_details_from_client_message
from litellm._uuid import uuid
from litellm.exceptions import MidStreamFallbackError
from litellm.types.llms.openai import ResponseFailedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents
class _ResponseIdentity(BaseModel):
model_config = ConfigDict(frozen=True, from_attributes=True)
id: str | None = None
model: str | None = None
created_at: int | None = None
class _StreamEvent(BaseModel):
model_config = ConfigDict(frozen=True, from_attributes=True)
type: str | None = None
sequence_number: int | None = None
response: _ResponseIdentity | None = None
class _FailureDetails(BaseModel):
model_config = ConfigDict(frozen=True, from_attributes=True)
message: str | None = None
code: str | int | None = None
type: str | None = None
status_code: int | None = None
@field_validator("message", mode="before")
@classmethod
def normalize_message(cls, value: object) -> str | None:
return value if isinstance(value, str) else None
@field_validator("code", mode="before")
@classmethod
def normalize_code(cls, value: object) -> str | int | None:
return value if isinstance(value, (str, int)) and not isinstance(value, bool) else None
@field_validator("type", mode="before")
@classmethod
def normalize_type(cls, value: object) -> str | None:
return value if isinstance(value, str) else None
def _original_failure(exception: Exception) -> Exception:
current = exception # rebind-ok: the recursion gate requires iterative wrapper traversal
while isinstance(current, MidStreamFallbackError) and current.original_exception is not None:
current = current.original_exception
return current
def _failure_details(original: Exception) -> _FailureDetails:
mapped: Final = _FailureDetails.model_validate(original)
body: Final = getattr(original, "body", None)
if not isinstance(body, Mapping):
return mapped
upstream: Final = _FailureDetails.model_validate(body)
return _FailureDetails(
message=upstream.message or mapped.message,
code=upstream.code if upstream.code is not None else mapped.code,
type=upstream.type or mapped.type,
status_code=mapped.status_code,
)
_CLIENT_ERROR_CODES: Final = MappingProxyType(
{
int(HTTPStatus.UNAUTHORIZED): "authentication_error",
int(HTTPStatus.FORBIDDEN): "permission_error",
int(HTTPStatus.NOT_FOUND): "not_found_error",
int(HTTPStatus.REQUEST_TIMEOUT): "request_timeout",
int(HTTPStatus.TOO_MANY_REQUESTS): "rate_limit_exceeded",
}
)
def _status_error_code(status_code: int | None) -> str:
if status_code is None or not HTTPStatus.BAD_REQUEST <= status_code < HTTPStatus.INTERNAL_SERVER_ERROR:
return "server_error"
return _CLIENT_ERROR_CODES.get(status_code, "invalid_request_error")
def _response_error_code(details: _FailureDetails) -> str:
for value in (details.code, details.type):
if value == "insufficient_quota":
return "insufficient_quota"
if value in (429, "429") or isinstance(value, str) and value.startswith("rate_limit"):
return "rate_limit_exceeded"
if isinstance(details.code, str) and details.code and not details.code.isdecimal():
return details.code
return _status_error_code(details.status_code)
class ResponsesStreamErrorState:
def __init__(self) -> None:
self.response_id: str | None = None
self.model: str | None = None
self.created_at: int | None = None
self.sequence_number = -1
self.terminal_emitted = False
self._pending_event: _StreamEvent | None = None
def observe_chunk(self, chunk: object) -> None:
self._pending_event = _StreamEvent.model_validate(chunk) if isinstance(chunk, (BaseModel, Mapping)) else None
def mark_emitted(self, frame: str | bytes) -> str | bytes:
event: Final = self._pending_event
if event is None:
return frame
if event.sequence_number is not None:
self.sequence_number = max(self.sequence_number, event.sequence_number)
if event.response is not None:
self.response_id = event.response.id or self.response_id
self.model = event.response.model or self.model
if event.response.created_at is not None:
self.created_at = event.response.created_at
if event.type in ("response.completed", "response.failed", "response.incomplete"):
self.terminal_emitted = True
return frame
def format_failure(self, exception: Exception) -> str | None:
if self.terminal_emitted:
return None
original: Final = _original_failure(exception)
details: Final = _failure_details(original)
response: Final = ResponsesAPIResponse.model_validate(
MappingProxyType(
{
"id": self.response_id or f"resp_{uuid.uuid4().hex}",
"object": "response",
"created_at": self.created_at if self.created_at is not None else int(time.time()),
"model": self.model,
"status": "failed",
"output": (),
"error": MappingProxyType(
{
"code": _response_error_code(details),
"message": redact_internal_details_from_client_message(details.message or str(original)),
}
),
}
)
)
event: Final = ResponseFailedEvent.model_validate(
MappingProxyType(
{
"type": ResponsesAPIStreamEvents.RESPONSE_FAILED,
"response": response,
"sequence_number": self.sequence_number + 1,
}
)
)
payload: Final = event.model_dump_json(exclude_none=True)
self.terminal_emitted = True
return f"event: response.failed\ndata: {payload}\n\n"

8
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -325,6 +325,14 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str:
return f"spend:model_access_group:{access_group_name}"
def project_cache_key(project_id: str) -> str:
return f"project_id:{project_id}"
def project_spend_counter_key(project_id: str) -> str:
return f"spend:project:{project_id}"
#: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds
#: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch.
END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__"

249
litellm/proxy/db/db_spend_update_writer.py
View file

@ -15,7 +15,7 @@ import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, overload
from urllib.parse import quote, unquote
from typing_extensions import LiteralString, ReadOnly, TypedDict
@ -31,6 +31,7 @@ from litellm.constants import (
from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
DB_RETRY_SAFE_ERROR_TYPES,
BaseDailySpendTransaction,
DailyAgentSpendTransaction,
@ -46,6 +47,7 @@ from litellm.proxy._types import (
SpendUpdateQueueItem,
ToolDiscoveryQueueItem,
)
from litellm.proxy.common_utils.user_api_key_cache import project_cache_key
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
build_bulk_upsert,
@ -64,6 +66,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
WindowSpendTransaction,
WindowSpendUpdateQueue,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
from litellm.proxy.spend_tracking.compression_savings import (
extract_compression_saved_tokens,
@ -122,6 +125,7 @@ class _SpendBatch(Protocol):
litellm_teammembership: BatchTable
litellm_organizationtable: BatchTable
litellm_organizationmembership: BatchTable
litellm_projecttable: BatchTable
litellm_tagtable: BatchTable
litellm_agentstable: BatchTable
litellm_modelaccessgroupbudgettable: BatchTable
@ -145,6 +149,30 @@ class _SpendTransactionManager(Protocol):
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
_DailySpendTransactionT = TypeVar("_DailySpendTransactionT", bound=BaseDailySpendTransaction)
class _DailySpendCommit(Protocol[_DailySpendTransactionT]):
async def __call__(
self,
*,
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
daily_spend_transactions: dict[str, _DailySpendTransactionT],
) -> None: ...
_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"})
def _daily_spend_commit_failure_is_requeue_safe(e: Exception) -> bool:
if isinstance(e, DB_CONNECTION_ERROR_TYPES):
return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES)
sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e)
return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES
def _timed_request_duration_ms(
payload: dict | SpendLogsPayload,
request_status: Literal["success", "failure"],
@ -300,6 +328,7 @@ class DBSpendUpdateWriter:
start_time: datetime,
end_time: datetime,
response_cost: float | None,
project_id: str | None = None,
) -> bool:
"""Record the request's spend, answering whether its cost still needs charging.
@ -382,6 +411,7 @@ class DBSpendUpdateWriter:
hashed_token=hashed_token,
team_id=team_id,
org_id=org_id,
project_id=project_id,
end_user_id=end_user_id,
prisma_client=prisma_client,
litellm_proxy_budget_name=litellm_proxy_budget_name,
@ -678,6 +708,7 @@ class DBSpendUpdateWriter:
litellm_proxy_budget_name: str | None,
payload: SpendLogsPayload,
request_model_access_groups: Sequence[str] = (),
project_id: str | None = None,
):
"""
Runs all 13 spend-update helpers sequentially inside a single asyncio task.
@ -741,6 +772,18 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
try:
await self._update_project_db(
response_cost=response_cost,
project_id=project_id,
prisma_client=prisma_client,
)
except Exception: # noqa: BLE001 # a project enqueue failure must not skip the sibling spend writes
verbose_proxy_logger.debug(
"_batch_database_updates: _update_project_db failed: %s",
traceback.format_exc(),
)
try:
await self._update_tag_db(
response_cost=response_cost,
@ -1003,6 +1046,32 @@ class DBSpendUpdateWriter:
)
raise e
async def _update_project_db(
self,
response_cost: float | None,
project_id: str | None,
prisma_client: PrismaClient | None,
) -> None:
if project_id is None or prisma_client is None:
return
try:
await self.spend_update_queue.add_update(
update=SpendUpdateQueueItem(
entity_type=Litellm_EntityType.PROJECT,
entity_id=project_id,
response_cost=response_cost,
)
)
except Exception as e:
spend_log_error(
"Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s",
project_id,
response_cost,
str(e),
exc=e,
)
raise e
async def _update_agent_db(
self,
response_cost: float | None,
@ -1240,18 +1309,19 @@ class DBSpendUpdateWriter:
if db_spend_update_transactions is not None:
verbose_proxy_logger.info(
"Spend tracking - committing spend updates from Redis to DB: "
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, "
"agents=%d, model_access_groups=%d",
len(db_spend_update_transactions.get("key_list_transactions") or {}),
len(db_spend_update_transactions.get("user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_list_transactions") or {}),
len(db_spend_update_transactions.get("org_list_transactions") or {}),
len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
len(db_spend_update_transactions.get("org_member_list_transactions") or {}),
len(db_spend_update_transactions.get("tag_list_transactions") or {}),
len(db_spend_update_transactions.get("agent_list_transactions") or {}),
len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}),
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, "
"projects=%d, tags=%d, agents=%d, model_access_groups=%d",
len(db_spend_update_transactions.get("key_list_transactions") or ()),
len(db_spend_update_transactions.get("user_list_transactions") or ()),
len(db_spend_update_transactions.get("team_list_transactions") or ()),
len(db_spend_update_transactions.get("org_list_transactions") or ()),
len(db_spend_update_transactions.get("end_user_list_transactions") or ()),
len(db_spend_update_transactions.get("team_member_list_transactions") or ()),
len(db_spend_update_transactions.get("org_member_list_transactions") or ()),
len(db_spend_update_transactions.get("project_list_transactions") or ()),
len(db_spend_update_transactions.get("tag_list_transactions") or ()),
len(db_spend_update_transactions.get("agent_list_transactions") or ()),
len(db_spend_update_transactions.get("model_access_group_list_transactions") or ()),
)
await self._commit_spend_updates_to_db(
prisma_client=prisma_client,
@ -1328,6 +1398,36 @@ class DBSpendUpdateWriter:
cronjob_id=DB_SPEND_UPDATE_JOB_NAME,
)
async def _flush_daily_spend_queue(
self,
queue: DailySpendUpdateQueue,
entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"],
commit: _DailySpendCommit[_DailySpendTransactionT],
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
) -> None:
transactions: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions()
try:
await commit(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions),
)
except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush
if not transactions:
return
spend_log_error(
"Spend tracking - failed to commit daily %s spend updates. "
"Re-queued %d rows for retry on next tick. Error: %s",
entity_type,
len(transactions),
str(e),
exc=e,
)
await queue.add_update(transactions)
async def _commit_spend_updates_to_db_without_redis_buffer(
self,
prisma_client: PrismaClient,
@ -1356,74 +1456,59 @@ class DBSpendUpdateWriter:
################## Daily Spend Update Transactions ##################
# Aggregate all in memory daily spend transactions and commit to db
daily_spend_update_transactions: Final = cast(
dict[str, DailyUserSpendTransaction],
await self.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_user_spend(
await self._flush_daily_spend_queue(
queue=self.daily_spend_update_queue,
entity_type="user",
commit=DBSpendUpdateWriter.update_daily_user_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_spend_update_transactions,
)
################## Daily Team Spend Update Transactions ##################
# Aggregate all in memory daily team spend transactions and commit to db
daily_team_spend_update_transactions: Final = cast(
dict[str, DailyTeamSpendTransaction],
await self.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_team_spend(
await self._flush_daily_spend_queue(
queue=self.daily_team_spend_update_queue,
entity_type="team",
commit=DBSpendUpdateWriter.update_daily_team_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_team_spend_update_transactions,
)
################## Daily Organization Spend Update Transactions ##################
# Aggregate all in memory daily org spend transactions and commit to db
daily_org_spend_update_transactions: Final = cast(
dict[str, DailyOrganizationSpendTransaction],
await self.daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_org_spend(
await self._flush_daily_spend_queue(
queue=self.daily_org_spend_update_queue,
entity_type="org",
commit=DBSpendUpdateWriter.update_daily_org_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_org_spend_update_transactions,
)
# NOTE: Daily tag spend is committed by a separate scheduler job.
################## Daily End-User Spend Update Transactions ##################
# Aggregate all in memory daily end-user spend transactions and commit to db
daily_end_user_spend_update_transactions: Final = cast(
dict[str, DailyEndUserSpendTransaction],
await self.daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_end_user_spend(
await self._flush_daily_spend_queue(
queue=self.daily_end_user_spend_update_queue,
entity_type="end_user",
commit=DBSpendUpdateWriter.update_daily_end_user_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_end_user_spend_update_transactions,
)
################## Daily Agent Spend Update Transactions ##################
# Aggregate all in memory daily agent spend transactions and commit to db
daily_agent_spend_update_transactions: Final = cast(
dict[str, DailyAgentSpendTransaction],
await self.daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_agent_spend(
await self._flush_daily_spend_queue(
queue=self.daily_agent_spend_update_queue,
entity_type="agent",
commit=DBSpendUpdateWriter.update_daily_agent_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_agent_spend_update_transactions,
)
################## Budget Window Spend Update Transactions ##################
@ -1460,19 +1545,15 @@ class DBSpendUpdateWriter:
Commit only tag spend updates to database.
This is called by a separate scheduler job at a longer interval.
"""
daily_tag_spend_update_transactions: Final = cast(
dict[str, DailyTagSpendTransaction],
await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
await self._flush_daily_spend_queue(
queue=self.daily_tag_spend_update_queue,
entity_type="tag",
commit=DBSpendUpdateWriter.update_daily_tag_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
if daily_tag_spend_update_transactions:
await DBSpendUpdateWriter.update_daily_tag_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_tag_spend_update_transactions,
)
async def _commit_daily_tag_spend_to_db_with_redis(
self,
prisma_client: PrismaClient,
@ -1797,6 +1878,22 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE PROJECT TABLE ###
project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions")
await DBSpendUpdateWriter._update_entity_spend_in_db(
entity_name="Project",
transactions=project_list_transactions,
table_accessor="litellm_projecttable",
where_field="project_id",
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
await DBSpendUpdateWriter._invalidate_project_caches(
project_ids=tuple(project_list_transactions or ()),
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TAG TABLE ###
tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"]
await DBSpendUpdateWriter._update_entity_spend_in_db(
@ -1835,11 +1932,23 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
@staticmethod
async def _invalidate_project_caches(project_ids: Sequence[str], proxy_logging_obj: ProxyLogging | None) -> None:
if not project_ids or proxy_logging_obj is None:
return
user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache")
if user_api_key_cache is None:
return
for project_id in project_ids:
await user_api_key_cache.async_delete_cache(key=project_cache_key(project_id))
@staticmethod
async def _update_entity_spend_in_db(
entity_name: str,
transactions: dict[str, float] | None,
table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"],
table_accessor: Literal[
"litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable", "litellm_projecttable"
],
where_field: str,
n_retry_times: int,
prisma_client: PrismaClient,
@ -2031,13 +2140,25 @@ class DBSpendUpdateWriter:
sql, params = build_bulk_upsert(table=table, batch=merged_batch)
await prisma_client.db.execute_raw(sql, *params)
except Exception as batch_error:
# Log detailed error information for debugging batch upsert failures
# This helps diagnose issues like unique constraint violations
if _daily_spend_commit_failure_is_requeue_safe(batch_error):
spend_log_error(
"Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s",
entity_type,
table.name,
len(transactions_to_process),
str(batch_error),
exc=batch_error,
)
raise
for key in transactions_to_process:
daily_spend_transactions.pop(key, None)
spend_log_error(
"Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s",
"Spend tracking - dropped %d daily %s spend rows: the failed statement may have "
"applied or the database refused the data, so re-sending it is not safe. "
"Table: %s, Error: %s",
len(transactions_to_process),
entity_type,
table.name,
len(transactions_to_process),
str(batch_error),
exc=batch_error,
)

7
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -70,6 +70,7 @@ _SpendTransactionField: TypeAlias = Literal[
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -83,6 +84,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -418,6 +420,10 @@ class RedisUpdateBuffer:
Litellm_EntityType.ORGANIZATION_MEMBER,
db_spend_update_transactions.get("org_member_list_transactions"),
),
(
Litellm_EntityType.PROJECT,
db_spend_update_transactions.get("project_list_transactions"),
),
(
Litellm_EntityType.TAG,
db_spend_update_transactions.get("tag_list_transactions"),
@ -885,6 +891,7 @@ class RedisUpdateBuffer:
org_member_list_transactions=_merged_entity_transactions(
list_of_transactions, "org_member_list_transactions"
),
project_list_transactions=_merged_entity_transactions(list_of_transactions, "project_list_transactions"),
tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"),
agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"),
model_access_group_list_transactions=_merged_entity_transactions(

4
litellm/proxy/db/db_transaction_queue/spend_update_queue.py
View file

@ -138,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
team_member_list_transactions={},
org_list_transactions={},
org_member_list_transactions={},
project_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
model_access_group_list_transactions={},
@ -152,6 +153,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions",
Litellm_EntityType.PROJECT: "project_list_transactions",
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions",
@ -192,6 +194,8 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions["org_list_transactions"]
elif dict_key == "org_member_list_transactions":
transactions_dict = db_spend_update_transactions["org_member_list_transactions"]
elif dict_key == "project_list_transactions":
transactions_dict = db_spend_update_transactions["project_list_transactions"]
elif dict_key == "tag_list_transactions":
transactions_dict = db_spend_update_transactions["tag_list_transactions"]
elif dict_key == "agent_list_transactions":

18
litellm/proxy/db/exception_handler.py
View file

@ -1,6 +1,8 @@
from collections.abc import Awaitable, Callable, Iterator
from typing import Any, Final, TypeVar
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
@ -17,6 +19,8 @@ _TRANSIENT_DB_UNAVAILABLE_MESSAGE: Final = (
"Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly."
)
_DATABASE_ERROR_META: Final = TypeAdapter(dict[str, object])
def _exception_chain(e: BaseException) -> Iterator[BaseException]:
current = e # rebind-ok: advances one link per iteration of the bounded walk
@ -221,6 +225,20 @@ class PrismaDBExceptionHandler:
or "write conflict or a deadlock" in error_message
)
@staticmethod
def postgres_sqlstate(e: Exception) -> str | None:
"""The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None."""
import prisma
if not isinstance(e, _exception_types(prisma.errors.DataError)):
return None
try:
meta: Final = _DATABASE_ERROR_META.validate_python(getattr(e, "meta", None))
except ValidationError:
return None
code: Final = meta.get("code")
return code if isinstance(code, str) else None
@staticmethod
def is_read_only_transaction_error(e: Exception) -> bool:
"""True iff ``e`` is Postgres SQLSTATE 25006 surfaced through prisma: the

5
litellm/proxy/db/spend_counter_reseed.py
View file

@ -26,6 +26,7 @@ from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_lookup_gate import db_lookup_gate
from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
BudgetWindowSpendRepository,
EndUserRepository,
@ -77,6 +78,7 @@ class SpendCounterReseed:
spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend
spend:user:{user_id} -> LiteLLM_UserTable.spend
spend:org:{org_id} -> LiteLLM_OrganizationTable.spend
spend:project:{project_id} -> LiteLLM_ProjectTable.spend
End-user and tag spend counters intentionally do not reseed here. Their
auth paths already load the corresponding objects via get_end_user_object()
@ -157,6 +159,9 @@ class SpendCounterReseed:
row = await OrganizationRepository(prisma_client).table.find_unique(
where={"organization_id": org_id}
)
elif counter_key.startswith("spend:project:"):
project_id: Final = counter_key[len("spend:project:") :]
row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id})
else:
return None
except Exception:

6
litellm/proxy/hooks/batch_rate_limiter.py
View file

@ -19,7 +19,7 @@ Quick summary:
import json
from collections.abc import Callable, Iterable, Mapping, Sequence
from datetime import datetime
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias
@ -661,7 +661,9 @@ class _PROXY_BatchRateLimiter(CustomLogger):
) or self.parallel_request_limiter.window_size
reset_time: Final = now + window_size if window_start is None else window_start + window_size
retry_after: Final = max(0, int(reset_time - now))
reset_time_formatted: Final = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC")
reset_time_formatted: Final = datetime.fromtimestamp(reset_time, tz=timezone.utc).strftime(
"%Y-%m-%d %H:%M:%S UTC"
)
remaining_display: Final = max(0, status["limit_remaining"])
current_limit: Final = status["current_limit"]

6
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -13,7 +13,7 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen
from contextlib import asynccontextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime
from datetime import datetime, timezone
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -3124,7 +3124,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
now = self._get_current_time().timestamp()
reset_time = now + self.window_size
reset_time_formatted = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC")
reset_time_formatted = datetime.fromtimestamp(reset_time, tz=timezone.utc).strftime(
"%Y-%m-%d %H:%M:%S UTC"
)
remaining_display = max(0, status["limit_remaining"])
rate_limit_type = status["rate_limit_type"]

12
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -267,6 +267,7 @@ class _ProxyDBLogger(CustomLogger):
start_time=actual_start_time,
end_time=datetime.now(),
org_id=user_api_key_dict.org_id,
project_id=user_api_key_dict.project_id,
)
@log_db_metrics
@ -318,6 +319,11 @@ class _ProxyDBLogger(CustomLogger):
user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
project_id: Final = (
project_id_value
if isinstance(project_id_value := metadata.get("user_api_key_project_id"), str)
else None
)
key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None))
end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None)
sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
@ -368,6 +374,7 @@ class _ProxyDBLogger(CustomLogger):
budget_reservation=budget_reservation,
request_tags=tags,
model_access_groups=model_access_groups,
project_id=project_id,
)
if not charged:
return
@ -501,6 +508,8 @@ class _ProxyDBLogger(CustomLogger):
metadata["user_api_key_team_id"] = key_obj.team_id
if metadata.get("user_api_key_org_id") is None:
metadata["user_api_key_org_id"] = key_obj.org_id
if metadata.get("user_api_key_project_id") is None:
metadata["user_api_key_project_id"] = key_obj.project_id
except Exception:
verbose_proxy_logger.debug(
"Failed to enrich failure metadata with key info for api_key=%s",
@ -651,6 +660,7 @@ async def _update_database_and_spend_counters(
budget_reservation: dict | None,
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
) -> bool:
if budget_reservation is not None:
await _reconcile_budget_reservation_before_db_update(
@ -668,6 +678,7 @@ async def _update_database_and_spend_counters(
start_time=start_time,
end_time=end_time,
org_id=org_id,
project_id=project_id,
)
except Exception:
if budget_reservation is not None:
@ -698,6 +709,7 @@ async def _update_database_and_spend_counters(
tags=request_tags,
request_started_at=start_time,
model_access_groups=model_access_groups,
project_id=project_id,
)
except Exception:
if budget_reservation is not None:

1
litellm/proxy/litellm_pre_call_utils.py
View file

@ -2367,6 +2367,7 @@ async def add_litellm_data_to_request(
# OTel layer can compute pre-request latency, including on the failure
# path after the logging object is popped.
data[_metadata_variable_name]["litellm_received_at"] = getattr(request.state, "litellm_received_at", None)
data[_metadata_variable_name]["llm_api_timing_windows"] = ()
# OTEL Controls / Tracing
# Add the OTEL Parent Trace before sending it LiteLLM

17
litellm/proxy/management_endpoints/team_admin_field_permissions.py
View file

@ -1,4 +1,5 @@
"""Proxy-wide allow-list of team-settings fields a team admin may change on /team/update."""
"""Proxy-wide allow-list of what a team admin may do on the teams they administer: team-settings fields on
/team/update, plus the ``projects`` permission for /project/new and /project/update."""
from collections.abc import Mapping
from dataclasses import dataclass
@ -21,6 +22,10 @@ TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_field
# TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset({"tpm_limit", "rpm_limit", "max_budget"})
TEAM_ADMIN_PROJECTS_PERMISSION: Final = "projects"
SUPPORTED_TEAM_ADMIN_PERMISSIONS: Final[frozenset[str]] = SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS | {
TEAM_ADMIN_PROJECTS_PERMISSION
}
_FIELD_LIST: Final = TypeAdapter(list[str])
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
@ -67,17 +72,23 @@ def resolve_team_admin_editable_fields(
"%s must be a list of field names; ignoring %r", TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, raw
)
return frozenset()
unsupported: Final = configured - supported
unsupported: Final = configured - supported - SUPPORTED_TEAM_ADMIN_PERMISSIONS
if unsupported:
verbose_proxy_logger.warning(
"%s ignores unsupported field(s) %s; supported: %s",
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
sorted(unsupported),
sorted(supported),
sorted(supported | SUPPORTED_TEAM_ADMIN_PERMISSIONS),
)
return configured & supported
def team_admin_may_manage_projects(general_settings: Mapping[str, object]) -> bool:
return TEAM_ADMIN_PROJECTS_PERMISSION in resolve_team_admin_editable_fields(
general_settings, frozenset({TEAM_ADMIN_PROJECTS_PERMISSION})
)
def _as_object(value: object) -> Mapping[str, object]:
try:
return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value)

5
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -3335,6 +3335,7 @@ async def team_member_add(
```
"""
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
premium_user,
@ -3429,6 +3430,10 @@ async def team_member_add(
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
await evict_and_broadcast(
cache_keys=tuple(sorted(user.user_id for user in updated_users)),
user_api_key_cache=user_api_key_cache,
)
await _evict_created_membership_caches(
user_ids=(tm.user_id for tm in updated_team_memberships),
team_id=data.team_id,

Some files were not shown because too many files have changed in this diff Show more