Merge remote-tracking branch 'origin/main' into litellm_org_alias_from_team

This commit is contained in:
ryan 2026-09-19 08:48:37 +00:00
commit 84e56a60d2
191 changed files with 13265 additions and 1153 deletions

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' }}

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

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**

View file

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

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**

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)

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,
)

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",

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"

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

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
}

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 {

View file

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

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

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
));
}
}

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())
}
}

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};

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()
);
}
}

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())
}
}

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);
}
}

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
));
}
}

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"

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;

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]

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)));
}
}

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;

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.

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.

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

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"
]
}

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);
}
}

View file

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

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);
});
}
}

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 },

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",

View file

@ -787,6 +787,7 @@ class InternalServerError(openai.InternalServerError):
super().__init__(
self.message, response=self.response, body=body
) # Call the base class constructor with the parameters it needs
self.type = "internal_server_error"
def __str__(self):
_message = self.message

View file

@ -44,6 +44,9 @@ from litellm.types.integrations.custom_logger import (
from litellm.types.integrations.websearch_interception import (
AnthropicSearchQuery,
AnthropicServerToolUseBlock,
RichWebSearchInput,
SearchFailed,
SearchOutcome,
WebSearchInterceptionConfig,
)
from litellm.types.llms.anthropic import AnthropicThinkingParam
@ -332,16 +335,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 +351,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 +926,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 +945,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 +988,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 +1000,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 +1018,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
@ -1152,7 +1145,9 @@ class WebSearchInterceptionLogger(CustomLogger):
"""Execute litellm.asearch() and build a Responses API rerun patch."""
search_tasks: Final = [
(
self._execute_search(tool_call["input"]["query"], kwargs=kwargs)
self._execute_search(
tool_call["input"]["query"], kwargs=kwargs, rich=self._rich_search_input(tool_call["input"])
)
if isinstance(tool_call.get("input"), dict) and tool_call["input"].get("query")
else self._create_empty_search_result()
)
@ -1306,7 +1301,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 +1339,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 +1354,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
@ -1376,7 +1365,9 @@ class WebSearchInterceptionLogger(CustomLogger):
query = tool_call["input"].get("query")
if query:
verbose_logger.debug("WebSearchInterception: Queuing search for query='%s'", query)
search_tasks.append(self._execute_search(query, kwargs=kwargs))
search_tasks.append(
self._execute_search(query, kwargs=kwargs, rich=self._rich_search_input(tool_call["input"]))
)
else:
verbose_logger.debug("WebSearchInterception: Tool call %s has no query", tool_call["id"])
# Add empty result for tools without query
@ -1385,27 +1376,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,10 +1423,66 @@ 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)
@staticmethod
def _rich_search_input(tool_input: object) -> RichWebSearchInput | None:
"""
Extract the optional objective/search_queries pair from a tool input.
Returns None when the input carries neither, so callers can pass the
result straight through as ``_execute_search``'s ``rich`` argument.
"""
if not isinstance(tool_input, Mapping):
return None
objective = tool_input.get("objective")
valid_objective = objective if isinstance(objective, str) and objective.strip() else None
raw_queries = tool_input.get("search_queries")
valid_queries: list[str] | None = None # mutable-ok: matches litellm.asearch's list[str] query parameter
if isinstance(raw_queries, Sequence) and not isinstance(raw_queries, str):
queries = [q for q in raw_queries if isinstance(q, str) and q.strip()]
if queries:
# Providers cap multi-query requests (Parallel drops queries
# past the fifth); trim here so nothing is silently ignored.
valid_queries = queries[:5]
if valid_objective is not None and valid_queries is not None:
return {"objective": valid_objective, "search_queries": valid_queries}
if valid_objective is not None:
return {"objective": valid_objective}
if valid_queries is not None:
return {"search_queries": valid_queries}
return None
@staticmethod
def _provider_supports_rich_search(search_provider: str | None) -> bool:
"""Whether the provider's search config accepts objective + multi-query input."""
if not search_provider:
return False
try:
from litellm.utils import ProviderConfigManager
except ImportError:
return False
# SearchProviders is a str enum, so an unknown provider string simply
# misses the config map and returns None rather than raising.
config = ProviderConfigManager.get_provider_search_config(search_provider) # pyright: ignore[reportArgumentType] -- SearchProviders is a str enum, so the router's provider string hashes to the matching member; unknown strings miss the map and yield None
return config is not None and config.supports_rich_search_input()
async def _execute_search(
self, query: str, kwargs: Mapping[str, object] | None = None
self,
query: str,
kwargs: Mapping[str, object] | None = None,
rich: RichWebSearchInput | None = None,
) -> tuple[str, SearchResponse | None]:
"""
Execute a single web search using router's search tools.
@ -1510,13 +1540,24 @@ class WebSearchInterceptionLogger(CustomLogger):
for key, value in search_litellm_params.items()
if key != "search_provider" and value is not None
}
# Forward the model's richer shape (objective + keyword queries)
# only to providers whose search API takes it natively; everyone
# else keeps the single query string the model also provided.
query_arg: str | list[str] = query # mutable-ok: litellm.asearch declares query as str | list[str]
if rich and self._provider_supports_rich_search(search_provider):
rich_queries = rich.get("search_queries")
if rich_queries:
query_arg = rich_queries
rich_objective = rich.get("objective")
if rich_objective and "objective" not in search_kwargs:
search_kwargs["objective"] = rich_objective
result: Final = (
await litellm.asearch(
query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
query=query_arg, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
)
if search_metadata is None
else await litellm.asearch(
query=query,
query=query_arg,
search_provider=search_provider,
litellm_metadata=search_metadata,
**_NO_ASEARCH_NAMED,
@ -1721,18 +1762,21 @@ class WebSearchInterceptionLogger(CustomLogger):
for tool_call in tool_calls:
# Handle both Anthropic-style input and OpenAI-style function.arguments
query = None
tool_args: dict | None = None # mutable-ok: the tool call's own arguments dict
if "input" in tool_call and isinstance(tool_call["input"], dict):
query = tool_call["input"].get("query")
tool_args = tool_call["input"]
query = tool_args.get("query")
elif "function" in tool_call:
func = tool_call["function"]
if isinstance(func, dict):
args = func.get("arguments", {})
if isinstance(args, dict):
tool_args = args
query = args.get("query")
if query:
verbose_logger.debug("WebSearchInterception: Queuing search for query='%s'", query)
search_tasks.append(self._execute_search(query, kwargs=kwargs))
search_tasks.append(self._execute_search(query, kwargs=kwargs, rich=self._rich_search_input(tool_args)))
else:
verbose_logger.debug("WebSearchInterception: Tool call %s has no query", tool_call.get("id"))
# Add empty result for tools without query

View file

@ -11,6 +11,50 @@ from typing import Any, Final
from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME
_WEB_SEARCH_TOOL_DESCRIPTION: Final = (
"Search the web for information. Use this when you need current "
"information or answers to questions that require up-to-date data."
)
def _web_search_input_schema() -> dict[str, object]: # mutable-ok: plain-dict tool shape, as the get_* builders
"""
JSON schema for the web search tool's input, shared by every tool format.
``query`` stays required so providers and callers that only understand a
single query string keep working unchanged. ``objective`` and
``search_queries`` are optional richer inputs; they are forwarded only to
search providers that support them (see
``BaseSearchConfig.supports_rich_search_input``).
"""
return {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query to execute",
},
"objective": {
"type": "string",
"description": (
"Natural-language description of the goal behind the "
"search, including any source or freshness requirements."
),
},
"search_queries": {
"type": "array",
"items": {"type": "string"},
"description": (
"Two to five short keyword queries (3-6 words each) "
"covering different angles of the objective, e.g. varying "
"names, synonyms, or phrasings. Provide together with "
"objective for the best results."
),
},
},
"required": ["query"],
}
def get_litellm_web_search_tool() -> dict[str, object]:
"""
@ -33,20 +77,8 @@ def get_litellm_web_search_tool() -> dict[str, object]:
"""
return {
"name": LITELLM_WEB_SEARCH_TOOL_NAME,
"description": (
"Search the web for information. Use this when you need current "
"information or answers to questions that require up-to-date data."
),
"input_schema": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query to execute",
}
},
"required": ["query"],
},
"description": _WEB_SEARCH_TOOL_DESCRIPTION,
"input_schema": _web_search_input_schema(),
}
@ -65,20 +97,8 @@ def get_litellm_web_search_tool_openai() -> dict[str, object]:
"type": "function",
"function": {
"name": LITELLM_WEB_SEARCH_TOOL_NAME,
"description": (
"Search the web for information. Use this when you need current "
"information or answers to questions that require up-to-date data."
),
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query to execute",
}
},
"required": ["query"],
},
"description": _WEB_SEARCH_TOOL_DESCRIPTION,
"parameters": _web_search_input_schema(),
},
}
@ -98,20 +118,8 @@ def get_litellm_web_search_tool_responses() -> dict[str, object]:
return {
"type": "function",
"name": LITELLM_WEB_SEARCH_TOOL_NAME,
"description": (
"Search the web for information. Use this when you need current "
"information or answers to questions that require up-to-date data."
),
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query to execute",
}
},
"required": ["query"],
},
"description": _WEB_SEARCH_TOOL_DESCRIPTION,
"parameters": _web_search_input_schema(),
}

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:
"""

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,

View file

@ -30,7 +30,7 @@ def get_formatted_prompt(
if c["type"] == "text":
prompt += c["text"]
if "tool_calls" in message:
for tool_call in message["tool_calls"]:
for tool_call in message["tool_calls"] or ():
if "function" in tool_call:
function_arguments = tool_call["function"]["arguments"]
prompt += function_arguments

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

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:

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 = "",

View file

@ -1253,10 +1253,9 @@ class AnthropicMessagesHandler(BaseTranslation):
Process output streaming response by applying guardrails to text content.
Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far.
With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite
written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked);
a rewrite on a stream that never reported a ``stop_reason`` has no write-back and is reported as
undeliverable, so the pipeline executor discards it and releases the original chunks.
With ``deliver_ended_stream_rewrites``, a stream whose guardrail rewrote the text gets the rewrite
written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked),
whether or not the stream ever reported a ``stop_reason``.
"""
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -1312,7 +1311,11 @@ class AnthropicMessagesHandler(BaseTranslation):
and guardrailed_texts
and guardrailed_texts[0] != string_so_far
):
self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
self._write_ended_stream_text_rewrite(
responses_so_far,
guardrailed_texts[0],
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
if deliver_ended_stream_rewrites:
returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
self._write_ended_stream_tool_call_rewrites(
@ -1354,9 +1357,11 @@ class AnthropicMessagesHandler(BaseTranslation):
raise
unended_texts: Final = _guardrailed_inputs.get("texts")
if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
self._write_ended_stream_text_rewrite(
responses_so_far,
unended_texts[0],
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
return responses_so_far
def _prepare_request_data(
@ -1450,26 +1455,40 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
@staticmethod
@classmethod
def _write_ended_stream_text_rewrite(
cls,
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
rewritten_text: str,
guardrail_name: str,
) -> None:
"""Deliver an ended-stream guardrail text rewrite by rewriting the
buffered chunks in place: the first ``text_delta`` carries the full
rewritten text and every later one is blanked, leaving the surrounding
message and content-block framing untouched."""
message and content-block framing untouched. A buffer with no
``text_delta`` has nowhere to carry the rewrite, so the pipeline
executor discards it and releases the original chunks."""
def is_text_delta(event: Mapping[str, object]) -> bool:
delta: Final = event.get("delta")
return (
event.get("type") == "content_block_delta"
and isinstance(delta, Mapping)
and delta.get("type") == "text_delta"
)
if not any(is_text_delta(event) for item in responses_so_far for event in cls._iter_sse_events(item)):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
replacements: Final = chain((rewritten_text,), repeat(""))
def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
delta: Final = event.get("delta")
if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping):
return None
if delta.get("type") != "text_delta":
if not is_text_delta(event):
return None
return _SSEFieldRewrite("delta", "text", next(replacements))
AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
cls._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
@classmethod
def _write_ended_stream_tool_call_rewrites(

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(

View file

@ -384,7 +384,7 @@ class LiteLLMAnthropicMessagesAdapter:
cache_control: Final = (
source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None)
)
if cache_control and model and (self.is_anthropic_claude_model(model) or self.is_bedrock_arn_model(model)):
if cache_control and model and self.target_consumes_cache_control(model):
# TypedDict objects support dict operations at runtime
# Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432)
if isinstance(target, dict):
@ -677,6 +677,10 @@ class LiteLLMAnthropicMessagesAdapter:
model_lower: Final = model.lower()
return "arn:" in model_lower and ":bedrock:" in model_lower
@classmethod
def target_consumes_cache_control(cls, model: str) -> bool:
return cls.is_anthropic_claude_model(model) or cls.is_bedrock_arn_model(model) or "gemini" in model.lower()
@staticmethod
def translate_thinking_for_model(
thinking: AnthropicThinkingParam,

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,

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

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,

View file

@ -95,6 +95,18 @@ class BaseSearchConfig:
"""
return "Unknown Search Provider"
def supports_rich_search_input(self) -> bool:
"""
Whether this provider's search API accepts a natural-language
objective plus multiple keyword queries in one request.
Integrations that collect the richer shape (e.g. websearch
interception) forward ``query`` as a list plus an ``objective``
optional param to providers that return True; every other provider
keeps receiving the single query string.
"""
return False
def get_http_method(self) -> Literal["GET", "POST"]:
"""
Get HTTP method for search requests.

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:

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)

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,

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
),
},
},
)

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,

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:]

View file

@ -18,6 +18,7 @@ import json
import time
import uuid
from collections.abc import Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union, cast
from typing_extensions import NotRequired, ReadOnly, TypedDict
@ -651,10 +652,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""Ended-stream path: rebuild the full response, run the non-streaming
output guardrail against it, and (when opted in) write any text or
tool-call rewrite back across the buffered chunks."""
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response)
await self.process_output_response(
@ -666,20 +664,59 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
)
if not deliver_ended_stream_rewrites:
return
guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown"
await self._write_ended_stream_text_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_texts=pre_guardrail_texts,
guardrail_name=guardrail_name,
)
self._write_ended_stream_tool_call_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
guardrail_name=guardrail_name,
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
@staticmethod
def _rebuild_ended_stream_per_choice(
responses_so_far: Sequence["ModelResponseStream"],
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> "ModelResponse":
"""``stream_chunk_builder`` folds every choice of a stream into one index-0
choice, so the stream is rebuilt one choice index at a time (every chunk
kept, its choices narrowed to that index, so usage-only chunks still
count) and the rebuilt choices are stitched into one response, each
carrying the index the stream gave it."""
choice_indices: Final = tuple(
sorted(frozenset(choice.index for response in responses_so_far for choice in response.choices))
)
rebuilt_by_index: Final = tuple(
(
index,
cast(
ModelResponse,
stream_chunk_builder(
chunks=[ # mutable-ok: callee takes a list
OpenAIChatCompletionsHandler._narrowed_to_choice(response, index)
for response in responses_so_far
],
logging_obj=litellm_logging_obj,
),
),
)
for index in choice_indices
)
(_, base_response), *_ = rebuilt_by_index
stitched_choices: Final = [ # mutable-ok: choices is a List field; a tuple there breaks model_dump round-trips
rebuilt.choices[0].model_copy(update=MappingProxyType({"index": index}))
for index, rebuilt in rebuilt_by_index
]
return base_response.model_copy(update=MappingProxyType({"choices": stitched_choices}))
@staticmethod
def _narrowed_to_choice(response: "ModelResponseStream", index: int) -> "ModelResponseStream":
narrowed: Final = [choice for choice in response.choices if choice.index == index] # mutable-ok: List field
return response.model_copy(update=MappingProxyType({"choices": narrowed}))
def build_stream_error_items(
self,
exc: "HTTPException",
@ -1058,39 +1095,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrailed_response: "ModelResponse",
pre_guardrail_texts: tuple[str | None, ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail text rewrites back across the buffered
chunks: the full rewritten text lands in the choice's first
content-carrying chunk and the rest are blanked, the same shape the
in-flight write-back uses. Chunks carrying only finish_reason or usage
stay untouched. A rewrite on a stream carrying more than one distinct
choice index is reported as undeliverable, so the pipeline executor
discards it and releases the original chunks."""
chunks, one rewrite per rebuilt choice index: the full rewritten text
lands in that choice's first content-carrying chunk and the rest are
blanked, the same shape the in-flight write-back uses. Chunks carrying
only finish_reason or usage stay untouched."""
post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response)
changed: Final = tuple(
after
for before, after in zip(pre_guardrail_texts, post_guardrail_texts)
if before is not None and after is not None and after != before
rewrites_by_choice: Final = MappingProxyType(
{
choice.index: after
for choice, before, after in zip(
guardrailed_response.choices, pre_guardrail_texts, post_guardrail_texts
)
if before is not None and after is not None and after != before
}
)
if not changed:
if not rewrites_by_choice:
return
stream_choice_indices: Final = frozenset(
choice.index for response in responses_so_far for choice in response.choices
)
if len(stream_choice_indices) != 1:
# stream_chunk_builder collapses every choice into one index-0
# choice, so a rewrite of the rebuilt response cannot be attributed
# back to a single choice on an n>1 stream: report it undeliverable
# rather than deliver the rewrite on the wrong choice
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
target_choice_index: Final = next(iter(stream_choice_indices))
await self._apply_guardrail_responses_to_output_streaming(
responses=responses_so_far,
guardrailed_texts=list(changed), # mutable-ok: callee takes lists
task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists
guardrailed_texts=list(rewrites_by_choice.values()), # mutable-ok: callee takes lists
task_mappings=[(index, None) for index in rewrites_by_choice], # mutable-ok: callee takes lists
)
@staticmethod

View file

@ -209,6 +209,7 @@ _TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | f
_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS
)
_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
_OUTPUT_TEXT_EVENT_TYPES: Final = frozenset({"response.output_text.delta", "response.output_text.done"})
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call_output": "output", "message": "content"}
)
@ -832,9 +833,10 @@ class OpenAIResponsesHandler(BaseTranslation):
(``response.output_text.delta`` / ``.done``,
``response.content_part.done``, ``response.output_item.done``) are synced
to the rewritten envelope too, so a client reading deltas sees the
rewrite instead of the raw model output; a rewrite observed where no
write-back is possible is reported as undeliverable, so the pipeline
executor discards it and releases the original events.
rewrite instead of the raw model output; a stream with no envelope
gets its rewrite spread over the buffered text events, and a rewrite
observed where no write-back is possible is reported as undeliverable,
so the pipeline executor discards it and releases the original events.
"""
if not responses_so_far:
return responses_so_far
@ -958,10 +960,9 @@ class OpenAIResponsesHandler(BaseTranslation):
return responses_so_far
# ------------------------------------------------------------------ #
# Fallback: apply guardrail to the accumulated text string. #
# No structured write-back is possible here; guardrails that only #
# need to block/flag (not rewrite) still work correctly, and a #
# rewrite a caller expects delivered is reported undeliverable. #
# Fallback: apply guardrail to the accumulated text string. With no #
# envelope to rewrite, a rewrite a caller expects delivered is spread #
# over the buffered text events instead. #
# ------------------------------------------------------------------ #
string_so_far: Final = self.get_streaming_string_so_far(responses_so_far)
if string_so_far:
@ -979,11 +980,54 @@ class OpenAIResponsesHandler(BaseTranslation):
)
fallback_texts: Final = fallback_outputs.get("texts")
if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
self._spread_text_rewrite_over_stream_events(
stream_events=responses_so_far,
rewritten_text=fallback_texts[0],
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
return responses_so_far
def _spread_text_rewrite_over_stream_events(
self,
stream_events: Sequence[Any],
rewritten_text: str,
guardrail_name: str,
) -> None:
"""Deliver a text rewrite on a stream with no completed envelope by
spreading it over the text parts the guardrail scanned, in stream
order: the whole rewrite on the first part and every later part
blanked, through the same sync the envelope path uses. A scanned
event the sync cannot place (one that is not an ``output_text`` delta
or done, or lacks integer ``output_index`` / ``content_index``) makes
the rewrite undeliverable, so the pipeline executor discards it and
releases the original events."""
scanned_events: Final = tuple(
event
for event in stream_events
if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str)
)
scanned_positions: Final = tuple(
dict.fromkeys(
(stream_item_field(event, "output_index"), stream_item_field(event, "content_index"))
for event in scanned_events
)
)
placeable_positions: Final = tuple(
(output_index, content_index)
for output_index, content_index in scanned_positions
if isinstance(output_index, int) and isinstance(content_index, int)
)
if len(placeable_positions) != len(scanned_positions) or any(
stream_item_field(event, "type") not in _OUTPUT_TEXT_EVENT_TYPES for event in scanned_events
):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
self._sync_stream_events_with_rewrites(
stream_events=stream_events,
rewrites_by_position=MappingProxyType(dict(zip(placeable_positions, chain((rewritten_text,), repeat(""))))),
)
@staticmethod
def _write_event_field(event: object, field: str, value: str) -> None:
if isinstance(event, dict):

View file

@ -90,6 +90,11 @@ class ParallelAISearchConfig(BaseSearchConfig):
def ui_friendly_name() -> str:
return "Parallel AI"
def supports_rich_search_input(self) -> bool:
# The v1 search API takes `objective` + multiple `search_queries`
# natively; sending both is the documented best practice.
return True
def validate_environment(
self,
headers: dict,

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))

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]

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:

View file

@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works
import re
from collections.abc import Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Final, Literal
from litellm.types.llms.openai import AllMessageValues
@ -57,7 +59,7 @@ def extract_ttl_from_cached_messages(messages: list[AllMessageValues]) -> str |
messages: List of messages to extract TTL from
Returns:
Optional[str]: TTL string in format "3600s" or None if not found/invalid
Optional[str]: TTL normalized to Gemini's "<seconds>s" form, or None if not found/invalid
"""
for message in messages:
if not is_cached_message(message):
@ -79,40 +81,29 @@ def extract_ttl_from_cached_messages(messages: list[AllMessageValues]) -> str |
if cache_control.get("type") != "ephemeral":
continue
ttl = cache_control.get("ttl")
if ttl and _is_valid_ttl_format(ttl):
return str(ttl)
normalized_ttl = _normalize_ttl_to_seconds(cache_control.get("ttl"))
if normalized_ttl is not None:
return normalized_ttl
return None
def _is_valid_ttl_format(ttl: str) -> bool:
"""
Validate TTL format. Should be a string ending with 's' for seconds.
Examples: "3600s", "7200s", "1.5s"
_TTL_PATTERN: Final = re.compile(r"^([0-9]*\.?[0-9]+)([smh])$")
_TTL_UNIT_SECONDS: Final = MappingProxyType({"s": 1, "m": 60, "h": 3600})
_LAST_EXPIRY_GOOGLE_ACCEPTS: Final = datetime(9999, 12, 31, 23, 59, 59, tzinfo=timezone.utc)
Args:
ttl: TTL string to validate
Returns:
bool: True if valid format, False otherwise
"""
def _normalize_ttl_to_seconds(ttl: object) -> str | None:
if not isinstance(ttl, str):
return False
# TTL should end with 's' and contain a valid number before it
pattern: Final = r"^([0-9]*\.?[0-9]+)s$"
match: Final = re.match(pattern, ttl)
if not match:
return False
try:
# Ensure the numeric part is valid and positive
numeric_part: Final = float(match.group(1))
return numeric_part > 0
except ValueError:
return False
return None
match: Final = _TTL_PATTERN.match(ttl)
if match is None:
return None
seconds: Final = round(float(match.group(1)) * _TTL_UNIT_SECONDS[match.group(2)], 9)
longest_ttl: Final = (_LAST_EXPIRY_GOOGLE_ACCEPTS - datetime.now(timezone.utc)).total_seconds()
if not 0 < seconds <= longest_ttl:
return None
return f"{seconds:.9f}".rstrip("0").rstrip(".") + "s"
def separate_cached_messages(

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),
)

View file

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

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}"}

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

View file

@ -25925,6 +25925,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"prompt_cache_min_tokens": 2048,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -27895,6 +27896,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"prompt_cache_min_tokens": 2048,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -48905,7 +48907,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": {
@ -63459,6 +63462,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",

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,7 +14,6 @@ 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/
@ -68,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.

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

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"

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(

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

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

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,

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",

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": [

View file

@ -512,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",
@ -532,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.
@ -889,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",
@ -898,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
@ -2601,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",

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()

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
@ -1216,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(
@ -2414,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
@ -2438,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.
@ -2455,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.
@ -2577,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,

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:

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"]

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"]

View file

@ -235,6 +235,8 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger):
return None
formatted_prompt: Final = get_formatted_prompt(data=data, call_type=call_type)
if not formatted_prompt:
return None
is_prompt_attack = False
prompt_injection_system_prompt: Final = getattr(

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

View file

@ -137,7 +137,7 @@ from litellm.types.router import (
updateDeployment,
updateLiteLLMParams,
)
from litellm.types.utils import without_server_derived_pricing
from litellm.types.utils import echoed_cost_map_pricing_fields, without_server_derived_pricing
from litellm.utils import get_utc_datetime
if TYPE_CHECKING:
@ -876,7 +876,11 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
_raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True))
merged_model_name: Final = updated_patch.model_name or db_model.model_name
merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True)
merged_model_info: Final[dict[str, object]] = db_model.model_info.model_dump(exclude_none=True)
stored_model_info: Final = db_model.model_info.model_dump(exclude_none=True)
echoed_pricing: Final = echoed_cost_map_pricing_fields(stored_model_info)
merged_model_info: Final[dict[str, object]] = {
k: v for k, v in stored_model_info.items() if k not in echoed_pricing
}
# update litellm params
if updated_patch.litellm_params:

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)

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,

View file

@ -7,6 +7,7 @@ from collections.abc import MutableMapping
from typing import Any, Final
from fastapi import Request
from starlette.routing import get_route_path
from starlette.types import ASGIApp, Receive, Scope, Send
import litellm
@ -15,6 +16,12 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
# Cache the header name at module level to avoid repeated enum attribute access
_AUTHORIZATION_HEADER: Final = SpecialHeaders.openai_authorization.value # "Authorization"
_METRICS_MOUNT: Final = "/metrics"
def _is_metrics_route(scope: Scope) -> bool:
route_path: Final = get_route_path(scope)
return route_path == _METRICS_MOUNT or route_path.startswith(_METRICS_MOUNT + "/")
class PrometheusAuthMiddleware:
@ -36,7 +43,7 @@ class PrometheusAuthMiddleware:
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
# Fast path: only inspect HTTP requests; pass through websocket/lifespan immediately
if scope["type"] != "http" or "/metrics" not in scope.get("path", ""):
if scope["type"] != "http" or not _is_metrics_route(scope):
await self.app(scope, receive, send)
return

View file

@ -111,7 +111,6 @@ from litellm.proxy._types import (
PassThroughGenericEndpoint,
ProxyErrorTypes,
ProxyException,
RoleBasedPermissions,
SpecialModelNames,
SupportedDBObjectType,
TeamDefaultSettings,
@ -148,11 +147,15 @@ from litellm.router_utils.auto_router_tuning_baseline import (
from litellm.router_utils.routing_groups import parse_routing_groups
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import (
PRICING_OVERRIDES_KEY,
ModelResponse,
ModelResponseStream,
StreamingChoices,
TextCompletionResponse,
TokenCountResponse,
echoed_cost_map_pricing_fields,
is_server_derived_pricing_key,
pricing_override_fields,
)
from litellm.utils import cost_map_omits_token_price, load_credentials_from_list
@ -317,6 +320,7 @@ from litellm.proxy.analytics_endpoints.analytics_endpoints import (
router as analytics_router,
)
from litellm.proxy.auth.auth_checks import (
ROLE_BASED_PERMISSIONS_ADAPTER,
ExperimentalUIJWTToken,
can_key_call_resolved_model,
get_team_object,
@ -1348,8 +1352,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
user_api_key_cache=user_api_key_cache,
)
if prompt_injection_detection_obj is not None: # [TODO] - REFACTOR THIS
prompt_injection_detection_obj.update_environment(router=llm_router)
ProxyStartupEvent._attach_router_to_prompt_injection_detectors(llm_router=llm_router)
verbose_proxy_logger.debug("prisma_client: %s", prisma_client)
if prisma_client is not None and litellm.max_budget > 0:
@ -4867,6 +4870,16 @@ def _bind_general_settings_store(settings: SettingsStore) -> None:
general_settings = settings # pyright: ignore[reportAssignmentType] # legacy global accepts mappings
@lru_cache(maxsize=4096)
def _log_ignored_cost_map_copy(model_id: str, fields: tuple[str, ...]) -> None:
verbose_proxy_logger.warning(
"Deployment %s stores a copy of the cost map in model_info (%s); ignoring it so the deployment follows the "
"current cost map. Set the price on litellm_params to override the cost map on purpose.",
model_id,
", ".join(fields),
)
class ProxyConfig:
"""
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
@ -6307,9 +6320,7 @@ class ProxyConfig:
### RBAC ###
rbac_role_permissions: Final = general_settings.get("role_permissions", None)
if rbac_role_permissions is not None:
general_settings["role_permissions"] = [ # validate role permissions
RoleBasedPermissions(**role_permission) for role_permission in rbac_role_permissions
]
ROLE_BASED_PERMISSIONS_ADAPTER.validate_python(rbac_role_permissions)
### SSRF URL VALIDATION SETTINGS ###
_apply_ssrf_general_settings(general_settings)
@ -6694,7 +6705,12 @@ class ProxyConfig:
model.model_info["id"] = model.model_id
if "db_model" in model.model_info and model.model_info["db_model"] is False:
model.model_info["db_model"] = db_model
_model_info = RouterModelInfo(**model.model_info)
echoed_pricing: Final = echoed_cost_map_pricing_fields(model.model_info)
if echoed_pricing:
_log_ignored_cost_map_copy(str(model.model_info["id"]), echoed_pricing)
_model_info = RouterModelInfo(
**MappingProxyType({k: v for k, v in model.model_info.items() if k not in echoed_pricing})
)
else:
_model_info = RouterModelInfo(id=model.model_id, db_model=db_model)
@ -9384,6 +9400,15 @@ def select_data_generator(
)
def _pricing_override_stamps(
model_info: Mapping[str, object], litellm_params: Mapping[str, object]
) -> Mapping[str, object]:
own_pricing: Final = MappingProxyType(
{k: v for k, v in litellm_params.items() if v is not None and is_server_derived_pricing_key(k)}
)
return MappingProxyType({**own_pricing, PRICING_OVERRIDES_KEY: pricing_override_fields(model_info, own_pricing)})
def get_litellm_model_info(model: dict = {}):
model_info: Final = model.get("model_info", {})
model_to_lookup = model.get("litellm_params", {}).get("model", None)
@ -9420,6 +9445,14 @@ def giveup(e):
class ProxyStartupEvent:
@staticmethod
def _attach_router_to_prompt_injection_detectors(llm_router: Router | None) -> None:
for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(
_OPTIONAL_PromptInjectionDetection
):
if isinstance(callback, _OPTIONAL_PromptInjectionDetection):
callback.update_environment(router=llm_router)
@staticmethod
async def refresh_model_info() -> None:
if llm_router is not None:
@ -13724,10 +13757,17 @@ def _enrich_model_info_with_litellm_data(
llm_router.get_discovered_model_info(model_info.get("id")) if llm_router is not None else MappingProxyType({})
)
unpriced: Final = cost_map_omits_token_price(model_info.get("id"), litellm_model_info.get("key"))
for k, v in MappingProxyType({**litellm_model_info, **discovered_model_info}).items():
if k not in model_info or (model_info[k] is None and k in discovered_model_info):
model_info[k] = None if unpriced and k in ("input_cost_per_token", "output_cost_per_token") else v
model["model_info"] = model_info
stamped_model_info: Final = MappingProxyType(
{**model_info, **_pricing_override_stamps(model_info, model.get("litellm_params") or MappingProxyType({}))}
)
model["model_info"] = {
**stamped_model_info,
**{
k: None if unpriced and k in ("input_cost_per_token", "output_cost_per_token") else v
for k, v in MappingProxyType({**litellm_model_info, **discovered_model_info}).items()
if k not in stamped_model_info or (stamped_model_info[k] is None and k in discovered_model_info)
},
}
# don't return the api key / vertex credentials
# don't return the llm credentials
model = remove_sensitive_info_from_deployment(model, excluded_keys={"litellm_credential_name"})

View file

@ -1,7 +1,8 @@
import asyncio
import contextlib
import json
import time
from collections.abc import AsyncIterator, Awaitable, Mapping
from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence
from enum import Enum
from functools import partial
from types import MappingProxyType
@ -12,10 +13,12 @@ import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from openai.types.responses.response_create_params import ResponseInputParam
from pydantic import BaseModel, ConfigDict, ValidationError
from starlette.websockets import WebSocket, WebSocketDisconnect
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import EMPTY_MAPPING
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_responses_api_usage as _blocked_responses_api_usage,
@ -1291,7 +1294,8 @@ async def cancel_response(
async def _read_ws_model_from_first_frame(
websocket: WebSocket,
) -> tuple | None:
query_model: str | None = None,
) -> tuple[str, str] | None:
"""Read the first WS frame and return (model, raw_message), or None on error.
Sends an appropriate error frame and closes the socket before returning None.
@ -1340,7 +1344,7 @@ async def _read_ws_model_from_first_frame(
await websocket.close(code=1008, reason="Invalid first message")
return None
model: Final = _extract_model_from_first_ws_event(first_event)
model: Final = query_model or _extract_model_from_first_ws_event(first_event)
if not model:
await websocket.send_text(
json.dumps(
@ -1371,6 +1375,38 @@ def _extract_model_from_first_ws_event(first_event: Any) -> str | None:
return (nested.get("model") if isinstance(nested, dict) else None) or first_event.get("model")
class _ResponseCreateRoutingHints(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
input: str | Sequence[object] | None = None
previous_response_id: str | None = None
response: "_ResponseCreateRoutingHints | None" = None
def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, object]:
try:
frame: Final = _ResponseCreateRoutingHints.model_validate_json(first_message)
except ValidationError:
return EMPTY_MAPPING
nested: Final = frame.response or frame
hints: Final = {
"input": frame.input if nested.input is None else nested.input,
"previous_response_id": (
frame.previous_response_id if nested.previous_response_id is None else nested.previous_response_id
),
}
return MappingProxyType({key: value for key, value in hints.items() if value is not None})
def _responses_ws_failure_frame(failure: Exception) -> str:
raw_status: Final = getattr(failure, "status_code", None)
status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500
error_type: Final = (
"rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error"
)
return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}})
async def _enforce_responses_ws_first_frame_model_auth(
request: Request,
model: str,
@ -1457,19 +1493,16 @@ async def responses_websocket_endpoint(
accept_kwargs["subprotocol"] = requested_protocols[0]
await websocket.accept(**accept_kwargs)
first_message: str | None = None
if not model:
result: Final = await _read_ws_model_from_first_frame(websocket)
if result is None:
return
model, first_message = result
result: Final = await _read_ws_model_from_first_frame(websocket, query_model=model)
if result is None:
return
resolved_model, first_message = result
data: dict[str, object] = {
"model": model,
"model": resolved_model,
"websocket": websocket,
"first_message": first_message,
}
if first_message is not None:
data["first_message"] = first_message
# Construct a synthetic Request for pre-call processing
headers_list: Final = list(websocket.scope.get("headers") or [])
@ -1482,7 +1515,7 @@ async def responses_websocket_endpoint(
request: Final = Request(scope=scope)
request._url = websocket.url
_body_bytes: Final = json.dumps({"model": model}).encode()
_body_bytes: Final = json.dumps({"model": resolved_model}).encode()
async def return_body():
return _body_bytes
@ -1492,10 +1525,10 @@ async def responses_websocket_endpoint(
# Phase 1: pre-call processing (auth, guardrails, rate limits)
base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:
if first_message is not None:
if not model:
await _enforce_responses_ws_first_frame_model_auth(
request=request,
model=model,
model=resolved_model,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
@ -1514,7 +1547,7 @@ async def responses_websocket_endpoint(
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=model,
model=resolved_model,
route_type="_aresponses_websocket",
)
except Exception as e:
@ -1536,16 +1569,31 @@ async def responses_websocket_endpoint(
await websocket.close(code=1008, reason="Pre-call error")
return
routed_data: Final = dict(
data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message)
)
# Phase 2: route to upstream provider
try:
data["user_api_key_dict"] = user_api_key_dict
llm_call: Final = await route_request(
data=data,
data=routed_data,
route_type="_aresponses_websocket",
llm_router=llm_router,
user_model=user_model,
)
await llm_call
except Exception:
failure: Final = await llm_call
if isinstance(failure, Exception):
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=failure,
request_data=routed_data,
)
except Exception as e:
verbose_proxy_logger.exception("Responses WebSocket error")
with contextlib.suppress(Exception):
await websocket.send_text(_responses_ws_failure_frame(e))
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=routed_data,
)
await websocket.close(code=1011, reason="Internal server error")

View file

@ -30,7 +30,7 @@ from litellm.proxy.config_resolvers.sso import (
resolve_sso_config,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
SUPPORTED_TEAM_ADMIN_PERMISSIONS,
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
@ -216,7 +216,7 @@ class UIThemeSettingsResponse(SettingsResponse):
"""Response model for UI theme settings"""
_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS))
_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_PERMISSIONS))
class UISettings(BaseModel):
@ -315,7 +315,8 @@ class UISettings(BaseModel):
default=(),
description=(
"Team settings fields a team admin may change on the teams they administer. "
"Empty means team admins cannot edit team settings at all. "
"Include 'projects' to let team admins create and update projects for those teams. "
"Empty means team admins cannot edit team settings or manage projects at all. "
"Proxy admins and org admins are not affected."
),
json_schema_extra={ # mutable-ok: pydantic only merges json_schema_extra when it is a plain dict
@ -1626,7 +1627,7 @@ async def update_ui_settings(
raise HTTPException(status_code=422, detail=e.errors())
unsupported_team_fields: Final = sorted(
frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_PERMISSIONS
)
if unsupported_team_fields:
raise HTTPException(
@ -1634,7 +1635,7 @@ async def update_ui_settings(
detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization
"error": (
f"{TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING} does not support {unsupported_team_fields}. "
f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)}."
f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_PERMISSIONS)}."
)
},
)

View file

@ -27,6 +27,7 @@ from datetime import date, datetime, timedelta, timezone
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from functools import partial
from itertools import takewhile
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -1009,6 +1010,7 @@ class _CallbackCapabilities:
has_guardrail: bool = False
has_pre_call_override: bool = False
has_content_enforcer: bool = False
has_moderation_override: bool = False
# Tuple[(resolved_callback, "override" | "apply_guardrail"), ...]
# Ordered the same as ``litellm.callbacks``; used to build the streaming
# iterator chain without re-scanning per request.
@ -1019,6 +1021,11 @@ class _CallbackCapabilities:
resolved_callbacks: tuple[object, ...] = field(default_factory=tuple)
def _overrides_moderation_hook(callback: CustomLogger) -> bool:
leaf_to_base: Final = takewhile(lambda klass: klass is not CustomLogger, type(callback).__mro__)
return any("async_moderation_hook" in klass.__dict__ for klass in leaf_to_base)
class ProxyLogging:
"""
Logging/Custom Handlers for proxy.
@ -2605,6 +2612,7 @@ class ProxyLogging:
has_guardrail = False
has_pre_call_override = False
has_content_enforcer = False
has_moderation_override = False
iterator_overrides: Final[list[tuple[Any, str]]] = [] # (callback, kind)
resolved_callbacks: Final[list[CustomLogger]] = []
@ -2623,6 +2631,8 @@ class ProxyLogging:
continue
if isinstance(resolved, CustomGuardrail):
has_guardrail = True
elif _overrides_moderation_hook(resolved):
has_moderation_override = True
# Use the same leaf-class ``__dict__`` check as the other hook
# capabilities: only callbacks that actually override the hook
# contribute to the flag. Setting this for every ``CustomLogger``
@ -2667,6 +2677,7 @@ class ProxyLogging:
has_guardrail=has_guardrail,
has_pre_call_override=has_pre_call_override,
has_content_enforcer=has_content_enforcer,
has_moderation_override=has_moderation_override,
iterator_overrides=tuple(iterator_overrides),
resolved_callbacks=tuple(resolved_callbacks),
)
@ -2728,20 +2739,27 @@ class ProxyLogging:
user_api_key_dict: UserAPIKeyAuth | None,
call_type: CallTypesLiteral,
):
"""
Runs the CustomGuardrail's async_moderation_hook() in parallel
"""
# Fast path: skip the entire guardrail scan when no CustomGuardrail
# callbacks are registered. Saves per-request iteration over
# ``litellm.callbacks`` plus an ``asyncio.gather([])`` round trip on
# deployments with no guardrails configured.
if not ProxyLogging._callback_capabilities().has_guardrail:
caps: Final = ProxyLogging._callback_capabilities()
if not caps.has_guardrail and not caps.has_moderation_override:
return data
# Step 1: Collect all guardrail tasks to run in parallel
guardrail_tasks: Final = []
for callback in litellm.callbacks:
if isinstance(callback, CustomGuardrail):
if (
isinstance(callback, CustomLogger)
and not isinstance(callback, CustomGuardrail)
and _overrides_moderation_hook(callback)
and user_api_key_dict is not None
):
guardrail_tasks.append(
callback.async_moderation_hook(
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
)
elif isinstance(callback, CustomGuardrail):
################################################################
# Check if guardrail should be run for GuardrailEventHooks.during_call hook
################################################################
@ -2749,7 +2767,7 @@ class ProxyLogging:
# V1 implementation - backwards compatibility
if callback.event_hook is None and hasattr(callback, "moderation_check"):
if callback.moderation_check == "pre_call":
return
continue
else:
# Main - V2 Guardrails implementation
from litellm.types.guardrails import GuardrailEventHooks

View file

@ -38,7 +38,8 @@ from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_pr
from ..llms.bedrock.realtime.handler import BedrockRealtime
from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context
from ..llms.openai.realtime.handler import OpenAIRealtime
from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig
from ..llms.vertex_ai.audio_transcription.realtime_transformation import is_vertex_speech_to_text_model
from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig, vertex_realtime_config
from ..llms.vertex_ai.vertex_llm_base import VertexBase
from ..llms.xai.realtime.handler import XAIRealtime
from ..utils import client as wrapper_client
@ -541,8 +542,6 @@ async def _arealtime(
or get_secret_str("VERTEXAI_LOCATION")
)
resolved_location: Final = vertex_llm_base.get_vertex_region(vertex_region=vertex_location, model=model)
(
access_token,
resolved_project,
@ -553,17 +552,28 @@ async def _arealtime(
timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS,
)
vertex_realtime_config: Final = VertexAIRealtimeConfig(
async def resolve_vertex_access_token() -> str:
refreshed_token, _ = await _resolve_vertex_access_token_bounded(
credentials=vertex_credentials,
project_id=resolved_project,
resolver=vertex_access_token_resolver,
timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS,
)
return refreshed_token
vertex_provider_config: Final = vertex_realtime_config(
model,
access_token=access_token,
resolve_access_token=resolve_vertex_access_token,
project=resolved_project,
location=resolved_location,
location=vertex_location,
)
await base_llm_http_handler.async_realtime(
model=model,
websocket=websocket,
logging_obj=litellm_logging_obj,
provider_config=vertex_realtime_config,
provider_config=vertex_provider_config,
api_base=dynamic_api_base or litellm_params.api_base,
api_key=None,
client=client,
@ -684,6 +694,11 @@ async def _realtime_health_check(
api_base=resolved_api_base or "https://api.x.ai/v1", query_params={"model": model}
)
elif custom_llm_provider == "vertex_ai":
if is_vertex_speech_to_text_model(model):
raise ValueError(
f"Realtime health checks are not supported for Speech-to-Text streaming model {model};"
" health check it with mode audio_transcription"
)
vertex_model_params: Final = dict(resolved_params)
resolved_location: Final = vertex_llm_base.get_vertex_region(
vertex_region=VertexBase.safe_get_vertex_ai_location(vertex_model_params),

View file

@ -22,7 +22,6 @@ from litellm.types.llms.openai import (
ContentPartAddedEvent,
ContentPartDoneEvent,
ContentPartDonePartOutputText,
ContentPartDonePartReasoningText,
FunctionCallArgumentsDeltaEvent,
FunctionCallArgumentsDoneEvent,
OutputItemAddedEvent,
@ -102,6 +101,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.sent_response_created_event: bool = False
self.sent_response_in_progress_event: bool = False
self.sent_output_item_added_event: bool = False
self.sent_message_item_added_event: bool = False
self.sent_content_part_added_event: bool = False
self.sent_output_text_done_event: bool = False
self.sent_output_content_part_done_event: bool = False
@ -111,6 +111,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.completed_response = None
self.final_text: str = ""
self._cached_item_id: str | None = None
self._message_output_index: int = 0
self._cached_response_id: str | None = None
self._buffered_chunk: ModelResponseStream | None = None
self._upstream_exhausted: bool = False
@ -563,7 +564,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._sequence_number += 1
event: Final = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=0,
output_index=self._message_output_index,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_item_id,
@ -585,13 +586,41 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
event: Final = ContentPartAddedEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
item_id=self._cached_item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
part=BaseLiteLLMOpenAIResponseObject(**{"type": "output_text", "text": "", "annotations": []}),
)
event.__dict__["sequence_number"] = self._sequence_number
return event
def _queue_message_item_added_events(self) -> None:
if self._cached_item_id is None:
self._cached_item_id = f"msg_{uuid.uuid4()}"
self.sent_message_item_added_event = True
self.sent_content_part_added_event = True
if self._cached_reasoning_item_id is not None:
self._message_output_index = self._next_tool_output_index
self._next_tool_output_index += 1
else:
self._message_output_index = 0
self._sequence_number += 1
event: Final = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=self._message_output_index,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_item_id,
"type": "message",
"role": "assistant",
"status": "in_progress",
"content": [],
}
),
)
event.__dict__["sequence_number"] = self._sequence_number
self._pending_response_events.append(event)
self._pending_response_events.append(self.create_content_part_added_event())
def _merge_provider_specific_fields(self, src: dict) -> None:
"""Merge provider_specific_fields using last-value-wins for lists.
@ -711,7 +740,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return OutputTextDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id=self._cached_item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
text=getattr(litellm_complete_object.choices[0].message, "content", "") or "",
)
@ -721,33 +750,24 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._cached_item_id = f"msg_{uuid.uuid4()}"
text: Final = getattr(litellm_complete_object.choices[0].message, "content", "") or ""
reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or ""
annotations: Final = getattr(litellm_complete_object.choices[0].message, "annotations", None)
part: PART_UNION_TYPES | None = None
if reasoning_content:
part = ContentPartDonePartReasoningText(
type="reasoning_text",
reasoning=reasoning_content,
)
else:
response_annotations: Final = (
LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations(
annotations=annotations
)
)
part = ContentPartDonePartOutputText(
type="output_text",
text=text,
annotations=response_annotations,
logprobs=None,
response_annotations: Final = (
LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations(
annotations=annotations
)
)
part: Final[PART_UNION_TYPES] = ContentPartDonePartOutputText(
type="output_text",
text=text,
annotations=response_annotations,
logprobs=None,
)
return ContentPartDoneEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
item_id=self._cached_item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
part=part,
)
@ -766,7 +786,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
)
return OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=0,
output_index=self._message_output_index,
sequence_number=1,
item=BaseLiteLLMOpenAIResponseObject(
**{
@ -832,6 +852,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def return_default_done_events(
self, litellm_complete_object: ModelResponse
) -> BaseLiteLLMOpenAIResponseObject | None:
if self.sent_message_item_added_event is False:
final_content: Final = litellm_complete_object.choices[0].message.content or ""
if not final_content:
self.sent_output_text_done_event = True
self.sent_output_content_part_done_event = True
self.sent_output_item_done_event = True
return None
self._queue_message_item_added_events()
return self._pending_response_events.pop(0)
if self.sent_output_text_done_event is False:
self.sent_output_text_done_event = True
return self.create_output_text_done_event(litellm_complete_object)
@ -936,31 +965,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return
# Default: message
self._cached_item_id = self._cached_item_id or f"msg_{uuid.uuid4()}"
event = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=0,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_item_id,
"type": "message",
"role": "assistant",
"status": "in_progress",
"content": [],
}
),
)
event.__dict__["sequence_number"] = self._sequence_number
self._pending_response_events.append(event)
# Emit content_part.added immediately after output_item.added for message
# items. The OpenAI Responses spec requires this event before any
# output_text.delta events so downstream parsers can initialize the
# text part structure.
if not self.sent_content_part_added_event:
self.sent_content_part_added_event = True
content_part_event: Final = self.create_content_part_added_event()
self._pending_response_events.append(content_part_event)
self._queue_message_item_added_events()
return
async def __anext__(
@ -1115,12 +1120,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.collected_chat_completion_chunks.append(
self._snapshot_chunk_for_stream_chunk_builder(cast(ModelResponseStream, chunk))
)
# Emit any just-queued output_item event
if self._pending_response_events:
return self._pending_response_events.pop(0)
response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk)
if response_api_chunk:
return response_api_chunk
self._pending_response_events.append(response_api_chunk)
if self._pending_response_events:
return self._pending_response_events.pop(0)
# Otherwise, loop to next chunk
except StopIteration:
return self.common_done_event_logic(sync_mode=True)
@ -1162,7 +1166,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
event = OutputTextAnnotationAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
item_id=item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
annotation_index=idx,
annotation=annotation_dict,
@ -1189,11 +1193,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
# Priority 2: Handle text deltas
delta_content: Final = self._get_delta_string_from_streaming_choices(chunk.choices)
if delta_content:
if not self.sent_message_item_added_event:
self._queue_message_item_added_events()
self._sequence_number += 1
text_delta_event: Final = OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id=item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
delta=delta_content,
)

View file

@ -1,5 +1,6 @@
import asyncio
import contextvars
import json
from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
@ -8,7 +9,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast
import httpx
from pydantic import BaseModel, TypeAdapter
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import assert_never
import litellm
@ -2274,6 +2275,27 @@ def _deployment_reasoning_default(kwargs: Mapping[str, object]) -> Reasoning | d
return _JSON_OBJECT_ADAPTER.validate_python(reasoning_effort) if isinstance(reasoning_effort, Mapping) else None
_RESPONSES_WS_ROUTING_HINT_KEYS: Final = frozenset({"input", "previous_response_id"})
def _first_ws_frame_with_routed_input(first_message: str, routed_input: object) -> str:
try:
frame: Final = _JSON_OBJECT_ADAPTER.validate_json(first_message)
except ValidationError:
return first_message
if frame is None or routed_input is None:
return first_message
raw_nested: Final = frame.get("response")
nested: Final = _JSON_OBJECT_ADAPTER.validate_python(raw_nested) if isinstance(raw_nested, Mapping) else None
if nested is not None and nested.get("input") is not None:
if nested["input"] == routed_input:
return first_message
return json.dumps({**frame, "response": {**nested, "input": routed_input}})
if frame.get("input") == routed_input:
return first_message
return json.dumps({**frame, "input": routed_input})
def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults:
default_reasoning: Final = _deployment_reasoning_default(kwargs)
candidate_params: Final[dict[str, object]] = {
@ -2295,11 +2317,11 @@ async def _aresponses_websocket(
api_key: str | None = None,
timeout: float | None = None,
**kwargs,
):
) -> Exception | None:
"""
Private function to handle the Responses API WebSocket mode.
For PROXY use only.
For PROXY use only. Returns the provider failure that ended the connection, if any.
Resolves the LLM provider from ``model``, looks up the matching
``BaseResponsesAPIConfig``, and hands off to
@ -2364,10 +2386,14 @@ async def _aresponses_websocket(
"api_base",
"api_key",
"timeout",
"first_message",
*_RESPONSES_WS_ROUTING_HINT_KEYS,
}
remaining_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _explicit_keys}
deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS}
first_message: Final = kwargs.get("first_message")
await base_llm_http_handler.async_responses_websocket(
return await base_llm_http_handler.async_responses_websocket(
model=resolved_model,
websocket=websocket,
logging_obj=litellm_logging_obj,
@ -2375,9 +2401,14 @@ async def _aresponses_websocket(
api_base=resolved_api_base,
api_key=resolved_api_key,
timeout=timeout,
first_message=(
_first_ws_frame_with_routed_input(first_message, kwargs.get("input"))
if isinstance(first_message, str)
else None
),
user_api_key_dict=kwargs.get("user_api_key_dict"),
litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
custom_llm_provider=_custom_llm_provider,
request_defaults=_build_responses_websocket_request_defaults(kwargs),
request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs),
**remaining_kwargs,
)

View file

@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import copy
import json
import time
import traceback
@ -154,7 +155,7 @@ def _load_json_value(payload: str | bytes) -> object:
return json.loads(payload)
def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
def _model_id_from_metadata(litellm_metadata: Mapping[str, object] | None) -> str | None:
model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None
model_id: Final = model_info.get("id") if _is_json_object(model_info) else None
return model_id if isinstance(model_id, str) else None
@ -229,6 +230,29 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None
return next((status for status in map(_status_code_for_error_field, fields) if status is not None), 500)
def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_message, error_type, error_code = _error_event_fields(error_obj)
status_code: Final = _status_code_for_error_fields(error_type, error_code)
error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
provider_exception: Final = BaseLLMException(
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {error_body}}}",
body=error_body,
)
try:
return litellm.exception_type(
model=model,
custom_llm_provider=custom_llm_provider,
original_exception=provider_exception,
completion_kwargs={},
extra_kwargs={},
)
except Exception as mapped_exception:
return mapped_exception
def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool:
if isinstance(mapped_exception, litellm.ContentPolicyViolationError):
return True
@ -592,26 +616,7 @@ class BaseResponsesAPIStreamingIterator:
)
def _map_error_event_exception(self, error_obj: object) -> Exception:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_message, error_type, error_code = _error_event_fields(error_obj)
status_code: Final = _status_code_for_error_fields(error_type, error_code)
error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
provider_exception: Final = BaseLLMException(
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {error_body}}}",
body=error_body,
)
try:
return litellm.exception_type(
model=self.model or "",
custom_llm_provider=self.custom_llm_provider or "",
original_exception=provider_exception,
completion_kwargs={},
extra_kwargs={},
)
except Exception as mapped_exception:
return mapped_exception
return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "")
def _maybe_raise_for_error_event(self, result: object) -> None:
chunk_type: Final = getattr(result, "type", None)
@ -1695,6 +1700,65 @@ RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"})
_RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"})
_RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
def _ws_event_error(event: Mapping[str, object]) -> object:
if event.get("type") == "error":
return event.get("error")
response: Final = event.get("response")
return response.get("error") if _is_json_object(response) else None
def _restore_input_item_ids(items: Sequence[object]) -> Sequence[object]:
return ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(copy.deepcopy(list(items))) # pyright: ignore[reportPrivateUsage] # same restore the HTTP responses path runs
def _restored_container_fields(container: Mapping[str, object]) -> Mapping[str, object]:
input_items: Final = container.get("input")
previous_response_id: Final = container.get("previous_response_id")
restored: Final = {
"input": _restore_input_item_ids(input_items) if _is_json_array(input_items) else input_items,
"previous_response_id": (
ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(previous_response_id)
if isinstance(previous_response_id, str)
else previous_response_id
),
}
return MappingProxyType({key: value for key, value in restored.items() if value != container.get(key)})
def _restore_wrapped_ids_in_response_create(msg_obj: Mapping[str, object]) -> dict[str, object] | None:
nested: Final = msg_obj.get("response")
nested_fields: Final = _restored_container_fields(nested) if _is_json_object(nested) else EMPTY_MAPPING
top_fields: Final = _restored_container_fields(msg_obj)
if not nested_fields and not top_fields:
return None
restored_nested: Final = (
{"response": {**nested, **nested_fields}} if _is_json_object(nested) and nested_fields else EMPTY_MAPPING
)
return {**msg_obj, **top_fields, **restored_nested}
def _wrap_output_item_encrypted_content(
event_obj: Mapping[str, object], litellm_metadata: Mapping[str, object]
) -> dict[str, object] | None:
if not litellm_metadata.get("encrypted_content_affinity_enabled"):
return None
model_id: Final = _model_id_from_metadata(litellm_metadata)
item: Final = event_obj.get("item")
if model_id is None or not _is_json_object(item):
return None
encrypted_content: Final = item.get("encrypted_content")
if not isinstance(encrypted_content, str) or not encrypted_content:
return None
wrapped_content: Final = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies
encrypted_content=encrypted_content, model_id=model_id
)
return {**event_obj, "item": {**item, "encrypted_content": wrapped_content}}
class ResponsesWebSocketStreaming:
"""
@ -1721,6 +1785,7 @@ class ResponsesWebSocketStreaming:
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
authorized_model: str | None = None,
custom_llm_provider: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
):
self.websocket = websocket
@ -1728,6 +1793,9 @@ class ResponsesWebSocketStreaming:
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: dict[str, object] = request_data or {}
litellm_metadata: Final = self.request_data.get("litellm_metadata")
self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {}
self.custom_llm_provider: str | None = custom_llm_provider
self.messages: list[_MutableJsonObject] = []
self.input_messages: list[dict[str, object]] = []
self.first_message = first_message
@ -1796,13 +1864,65 @@ class ResponsesWebSocketStreaming:
if self.logging_obj:
self.logging_obj.pre_call(input=message, api_key="")
def _failure_exception(self) -> Exception | None:
failed_event: Final = next(
(event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None
)
if failed_event is None:
return None
return _map_stream_error_to_exception(
_ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or ""
)
async def _log_messages(self) -> None:
if not self.logging_obj:
return
if self.input_messages:
self.logging_obj.model_call_details["messages"] = self.input_messages
if self.messages:
if not self.messages:
return
exception: Final = self._failure_exception()
if exception is None:
asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True))
return
self._record_usage_for_failure()
traceback_exception: Final = "".join(traceback.format_exception(exception))
asyncio.create_task(
self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True)
)
def _record_usage_for_failure(self) -> None:
from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor
from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject
usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
self.messages
)
tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(self.messages)
service_tier: Final = next(iter(tier_partition)) if len(tier_partition) == 1 else None
logging_result: Final = LiteLLMRealtimeStreamLoggingObject(
usage=usage, results=self.messages, service_tier=service_tier
)
response_cost: Final = self.logging_obj._response_cost_calculator(result=logging_result) or 0.0 # pyright: ignore[reportPrivateUsage] # as the HTTP streaming iterator does
self.logging_obj.record_partial_usage_for_failure(usage, response_cost)
def _wrap_response_event(self, response_str: str) -> str:
try:
event_obj: Final = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
response: Final = event_obj.get("response")
if _is_json_object(response):
wrapped_response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies
responses_api_response=response,
custom_llm_provider=self.custom_llm_provider,
litellm_metadata=self.litellm_metadata,
)
return json.dumps({**event_obj, "response": wrapped_response})
if event_obj.get("type") not in _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES:
return response_str
wrapped_event: Final = _wrap_output_item_encrypted_content(event_obj, self.litellm_metadata)
return response_str if wrapped_event is None else json.dumps(wrapped_event)
async def backend_to_client(self) -> None:
"""Forward events from backend WebSocket to the client."""
@ -1839,12 +1959,13 @@ class ResponsesWebSocketStreaming:
unmasked_str = self._unmask_response_event(response_str)
output_masked_str = await self._mask_response_completed(unmasked_str)
wrapped_str = self._wrap_response_event(output_masked_str)
# Log the output-masked form so PII redacted by apply_to_output
# guardrails does not appear in success logs.
self._store_event(output_masked_str)
self._store_event(wrapped_str)
await self.websocket.send_text(output_masked_str)
await self.websocket.send_text(wrapped_str)
except websockets.exceptions.ConnectionClosed as e:
verbose_logger.debug("Responses WS backend connection closed: %s", e)
@ -1913,19 +2034,22 @@ class ResponsesWebSocketStreaming:
if parsed.get("type") != "response.create":
return message
msg_obj: Final = self._with_request_defaults(parsed)
defaults_applied: Final = msg_obj != parsed
authorized_obj: Final = self._with_request_defaults(parsed)
defaults_applied: Final = authorized_obj != parsed
# Always enforce the authorized model, even when PII masking is off.
model_modified: Final = self._enforce_authorized_model(msg_obj)
model_modified: Final = self._enforce_authorized_model(authorized_obj)
restored_obj: Final = _restore_wrapped_ids_in_response_create(authorized_obj)
msg_obj: Final = authorized_obj if restored_obj is None else restored_obj
frame_modified: Final = model_modified or restored_obj is not None or defaults_applied
if not self.guardrail_callbacks:
return json.dumps(msg_obj) if model_modified or defaults_applied else message
return json.dumps(msg_obj) if frame_modified else message
if "metadata" not in self.request_data:
self.request_data["metadata"] = {}
modified = model_modified or defaults_applied
modified = frame_modified
guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks)
for cb in guardrail_cbs:
presidio_config = cb.get_presidio_settings_from_request_data(self.request_data)
@ -2209,8 +2333,7 @@ class ResponsesWebSocketStreaming:
except Exception as e:
verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
async def bidirectional_forward(self) -> None:
"""Run both forwarding directions concurrently."""
async def bidirectional_forward(self) -> Exception | None:
forward_task: Final = asyncio.create_task(self.backend_to_client())
try:
await self.client_to_backend()
@ -2227,6 +2350,7 @@ class ResponsesWebSocketStreaming:
await self.backend_ws.close()
except Exception:
pass
return self._failure_exception()
# ---------------------------------------------------------------------------

View file

@ -0,0 +1,57 @@
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
@dataclass(frozen=True, slots=True)
class HttpSettings:
ssl_verify: bool | str
ssl_certificate: str | None
ssl_security_level: str | None
ssl_ecdh_curve: str | None
force_ipv4: bool
http2: bool
aiohttp_trust_env: bool
disable_aiohttp_trust_env: bool
disable_aiohttp_transport: bool
user_agent: str
@dataclass(frozen=True, slots=True)
class UrlPolicy:
user_url_validation: bool
user_url_allowed_hosts: Sequence[str]
def warn(message: str) -> None:
from litellm._logging import verbose_logger
verbose_logger.warning("%s", message)
def url_policy() -> UrlPolicy:
import litellm
return UrlPolicy(
user_url_validation=litellm.user_url_validation,
user_url_allowed_hosts=litellm.user_url_allowed_hosts,
)
def http_settings() -> HttpSettings:
import litellm
from litellm.llms.custom_httpx.http_handler import default_user_agent
return HttpSettings(
ssl_verify=litellm.ssl_verify,
ssl_certificate=litellm.ssl_certificate,
ssl_security_level=litellm.ssl_security_level,
ssl_ecdh_curve=litellm.ssl_ecdh_curve,
force_ipv4=litellm.force_ipv4,
http2=litellm.http2,
aiohttp_trust_env=litellm.aiohttp_trust_env,
disable_aiohttp_trust_env=litellm.disable_aiohttp_trust_env,
disable_aiohttp_transport=litellm.disable_aiohttp_transport,
user_agent=default_user_agent(),
)

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