mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge remote-tracking branch 'origin/main' into litellm_org_alias_from_team
This commit is contained in:
commit
84e56a60d2
191 changed files with 13265 additions and 1153 deletions
71
.github/workflows/issue_fixed_comment.yml
vendored
Normal file
71
.github/workflows/issue_fixed_comment.yml
vendored
Normal 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
130
AGENTS.md
|
|
@ -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
129
CLAUDE.md
|
|
@ -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
|
||||
|
|
@ -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**
|
||||
|
|
|
|||
|
|
@ -1 +1 @@
|
|||
Read @CLAUDE.md for coding guidelines
|
||||
Read @AGENTS.md for coding guidelines
|
||||
|
|
|
|||
|
|
@ -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**
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ Endpoints for /project operations
|
|||
#### PROJECT MANAGEMENT ####
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
|
|
@ -22,7 +22,11 @@ from litellm._uuid import uuid
|
|||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.auth_checks import delete_cached_project_object
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
|
||||
from litellm.proxy.management_endpoints.common_utils import (
|
||||
_is_user_team_admin, # pyright: ignore[reportPrivateUsage] # shared owner of team-admin membership
|
||||
_set_object_metadata_field,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
management_endpoint_wrapper,
|
||||
)
|
||||
|
|
@ -82,37 +86,38 @@ async def _check_user_permission_for_project(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
team_id: str | None,
|
||||
prisma_client: PrismaClient,
|
||||
general_settings: Mapping[str, object],
|
||||
require_admin: bool = False,
|
||||
team_object: LiteLLM_TeamTable | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if user has permission to manage a project.
|
||||
|
||||
Returns True if user is proxy admin or team admin (when team_id provided).
|
||||
Returns True if user is proxy admin, or a team admin of ``team_id`` when the
|
||||
``team_admin_editable_team_fields`` setting grants team admins the ``projects`` permission.
|
||||
If require_admin=True, only proxy admins are allowed.
|
||||
|
||||
If team_object is provided, it will be used instead of fetching from DB
|
||||
(avoids duplicate DB queries when team was already fetched for validation).
|
||||
"""
|
||||
is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
if require_admin:
|
||||
if require_admin or is_proxy_admin:
|
||||
return is_proxy_admin
|
||||
|
||||
if is_proxy_admin:
|
||||
return True
|
||||
|
||||
if not team_id or not user_api_key_dict.user_id:
|
||||
if not team_id or not user_api_key_dict.user_id or not team_admin_may_manage_projects(general_settings):
|
||||
return False
|
||||
|
||||
team = team_object
|
||||
if team is None:
|
||||
team = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
|
||||
team_row: Final = (
|
||||
team_object
|
||||
if team_object is not None
|
||||
else await _team_table(prisma_client).find_unique(where={"team_id": team_id})
|
||||
)
|
||||
if team_row is None:
|
||||
return False
|
||||
|
||||
if team and team.admins:
|
||||
return user_api_key_dict.user_id in team.admins
|
||||
|
||||
return False
|
||||
team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump())
|
||||
return _is_user_team_admin(user_api_key_dict, team) or user_api_key_dict.user_id in (team.admins or [])
|
||||
|
||||
|
||||
async def _validate_team_exists(
|
||||
|
|
@ -531,6 +536,7 @@ async def new_project(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
team_id=data.team_id,
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings,
|
||||
team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()),
|
||||
)
|
||||
|
||||
|
|
@ -735,6 +741,7 @@ async def update_project(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
team_id=existing_project.team_id,
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings,
|
||||
)
|
||||
|
||||
if not has_permission:
|
||||
|
|
@ -751,6 +758,7 @@ async def update_project(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
team_id=data.team_id,
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings,
|
||||
team_object=(
|
||||
LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()) if target_team_obj else None
|
||||
),
|
||||
|
|
@ -877,7 +885,7 @@ async def delete_project(
|
|||
}'
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
|
||||
from litellm.proxy.proxy_server import general_settings, premium_user, prisma_client, user_api_key_cache
|
||||
|
||||
try:
|
||||
if not premium_user:
|
||||
|
|
@ -899,6 +907,7 @@ async def delete_project(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
team_id=None,
|
||||
prisma_client=prisma_client,
|
||||
general_settings=general_settings,
|
||||
require_admin=True,
|
||||
)
|
||||
|
||||
|
|
|
|||
19
litellm-rust/Cargo.lock
generated
19
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,16 +1,21 @@
|
|||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use litellm_auth_gcp::VertexAuth;
|
||||
use litellm_host::{
|
||||
event::{CallEvent, MachineEvent, WireRequest},
|
||||
host::{Host, HostOp, HostResult},
|
||||
machine::{HostFailure, Machine, MachineStep},
|
||||
};
|
||||
use litellm_http::{HttpClientPool, HttpSettings, Resolution};
|
||||
use litellm_llms::{
|
||||
base_llm::ocr::{
|
||||
error::Error as OcrError,
|
||||
transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig},
|
||||
},
|
||||
custom_httpx::llm_http_handler::OcrClient,
|
||||
custom_httpx::{
|
||||
llm_http_handler::OcrClient,
|
||||
media::{PublicDnsResolver, UrlPolicy},
|
||||
},
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
|
@ -171,25 +176,24 @@ async fn facade_retains_native_response_when_requested() {
|
|||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn facade_uses_the_injected_http_client() {
|
||||
async fn ocr_client_uses_the_injected_http_pool_configuration() {
|
||||
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
|
||||
let mut default_headers = reqwest::header::HeaderMap::new();
|
||||
default_headers.insert(
|
||||
"x-transport-owner",
|
||||
reqwest::header::HeaderValue::from_static("host"),
|
||||
);
|
||||
let provider_http = reqwest::Client::builder()
|
||||
.default_headers(default_headers)
|
||||
.build()
|
||||
.unwrap();
|
||||
crate::ocr::client::perform(
|
||||
&OcrClient::new(provider_http).unwrap(),
|
||||
wire_request("mistral/model", &base, json!({})),
|
||||
let settings = HttpSettings {
|
||||
user_agent: Some("host-owned/1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let client = OcrClient::new(
|
||||
&HttpClientPool::new(Arc::new(PublicDnsResolver)),
|
||||
&Resolution::from(&settings).config,
|
||||
UrlPolicy::default(),
|
||||
VertexAuth::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({})))
|
||||
.await
|
||||
.unwrap();
|
||||
server.await.unwrap();
|
||||
assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host"));
|
||||
assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1"));
|
||||
}
|
||||
|
||||
fn event_name(event: &CallEvent) -> &'static str {
|
||||
|
|
|
|||
1
litellm-rust/crates/http/AGENTS.md
Normal file
1
litellm-rust/crates/http/AGENTS.md
Normal file
|
|
@ -0,0 +1 @@
|
|||
- https://github.com/BerriAI/litellm-docs/blob/main/docs/guides/security_settings.md
|
||||
18
litellm-rust/crates/http/Cargo.toml
Normal file
18
litellm-rust/crates/http/Cargo.toml
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
[package]
|
||||
name = "litellm-http"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
http.workspace = true
|
||||
hyper-util.workspace = true
|
||||
reqwest.workspace = true
|
||||
rustls.workspace = true
|
||||
thiserror.workspace = true
|
||||
webpki-roots.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
296
litellm-rust/crates/http/src/config.rs
Normal file
296
litellm-rust/crates/http/src/config.rs
Normal file
|
|
@ -0,0 +1,296 @@
|
|||
use std::{
|
||||
net::{IpAddr, Ipv4Addr},
|
||||
path::PathBuf,
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
error::Error,
|
||||
settings::{HttpSettings, SslVerify, TcpKeepalive},
|
||||
tls::{CipherSelection, KeyExchangeGroup, Tls12CipherSuite, Unsupported},
|
||||
};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum Verify {
|
||||
Disabled,
|
||||
CaBundle(PathBuf),
|
||||
BuiltInRoots,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct HttpClientConfig {
|
||||
pub verify: Verify,
|
||||
pub client_certificate: Option<PathBuf>,
|
||||
pub key_exchange_group: Option<KeyExchangeGroup>,
|
||||
pub tls12_cipher_suites: Option<Vec<Tls12CipherSuite>>,
|
||||
pub force_ipv4: bool,
|
||||
pub http2: bool,
|
||||
pub user_agent: Option<String>,
|
||||
pub trust_proxy_env: bool,
|
||||
pub connect_timeout: Duration,
|
||||
pub tcp_keepalive: Option<TcpKeepalive>,
|
||||
pub pool_idle_timeout: Duration,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Resolution {
|
||||
pub config: HttpClientConfig,
|
||||
pub unsupported: Vec<Unsupported>,
|
||||
}
|
||||
|
||||
impl From<&HttpSettings> for Verify {
|
||||
fn from(settings: &HttpSettings) -> Self {
|
||||
match &settings.ssl_verify {
|
||||
Some(SslVerify::Disabled) => Self::Disabled,
|
||||
Some(SslVerify::CaBundle(path)) => Self::CaBundle(path.clone()),
|
||||
Some(SslVerify::Enabled) | None => settings
|
||||
.ssl_cert_file
|
||||
.clone()
|
||||
.map_or(Self::BuiltInRoots, Self::CaBundle),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&HttpSettings> for Resolution {
|
||||
fn from(settings: &HttpSettings) -> Self {
|
||||
let curve = settings
|
||||
.ssl_ecdh_curve
|
||||
.as_deref()
|
||||
.map(str::parse::<KeyExchangeGroup>)
|
||||
.transpose();
|
||||
let ciphers = settings
|
||||
.ssl_security_level
|
||||
.as_deref()
|
||||
.map(CipherSelection::from)
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
config: HttpClientConfig {
|
||||
verify: Verify::from(settings),
|
||||
client_certificate: settings.ssl_certificate.clone(),
|
||||
key_exchange_group: curve.clone().ok().flatten(),
|
||||
tls12_cipher_suites: ciphers.tls12_cipher_suites,
|
||||
force_ipv4: settings.force_ipv4,
|
||||
http2: settings.http2,
|
||||
user_agent: settings.user_agent.clone(),
|
||||
trust_proxy_env: settings.trust_proxy_env,
|
||||
connect_timeout: settings.connect_timeout,
|
||||
tcp_keepalive: settings.tcp_keepalive,
|
||||
pool_idle_timeout: settings.pool_idle_timeout,
|
||||
},
|
||||
unsupported: curve.err().into_iter().chain(ciphers.unsupported).collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(config: &HttpClientConfig) -> Result<Self, Self::Error> {
|
||||
let base = reqwest::Client::builder()
|
||||
.use_preconfigured_tls(rustls::ClientConfig::try_from(config)?)
|
||||
.connect_timeout(config.connect_timeout)
|
||||
.pool_idle_timeout(config.pool_idle_timeout);
|
||||
let with_keepalive = match config.tcp_keepalive {
|
||||
None => base,
|
||||
Some(keepalive) => base
|
||||
.tcp_keepalive(keepalive.idle)
|
||||
.tcp_keepalive_interval(keepalive.interval)
|
||||
.tcp_keepalive_retries(keepalive.retries),
|
||||
};
|
||||
let with_address = if config.force_ipv4 {
|
||||
with_keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
|
||||
} else {
|
||||
with_keepalive
|
||||
};
|
||||
let with_protocol = if config.http2 {
|
||||
with_address
|
||||
} else {
|
||||
with_address.http1_only()
|
||||
};
|
||||
let with_agent = match &config.user_agent {
|
||||
Some(agent) => with_protocol.user_agent(agent),
|
||||
None => with_protocol,
|
||||
};
|
||||
Ok(if config.trust_proxy_env {
|
||||
with_agent
|
||||
} else {
|
||||
with_agent.no_proxy()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn settings(ssl_verify: Option<SslVerify>, ssl_cert_file: Option<&str>) -> HttpSettings {
|
||||
HttpSettings {
|
||||
ssl_verify,
|
||||
ssl_cert_file: ssl_cert_file.map(PathBuf::from),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::default(settings(None, None), Verify::BuiltInRoots)]
|
||||
#[case::setting_disables(
|
||||
settings(Some(SslVerify::Disabled), Some("/env/roots.pem")),
|
||||
Verify::Disabled
|
||||
)]
|
||||
#[case::setting_bundle(
|
||||
settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")),
|
||||
Verify::CaBundle("/configured.pem".into())
|
||||
)]
|
||||
#[case::enabled_uses_cert_file(
|
||||
settings(Some(SslVerify::Enabled), Some("/env/roots.pem")),
|
||||
Verify::CaBundle("/env/roots.pem".into())
|
||||
)]
|
||||
#[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), Verify::CaBundle("/env/roots.pem".into()))]
|
||||
fn verify_follows_setting_then_cert_file(
|
||||
#[case] settings: HttpSettings,
|
||||
#[case] expected: Verify,
|
||||
) {
|
||||
let config = Resolution::from(&settings).config;
|
||||
assert_eq!(config.verify, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::x25519("X25519", Some(KeyExchangeGroup::X25519))]
|
||||
#[case::openssl_p256("prime256v1", Some(KeyExchangeGroup::Secp256r1))]
|
||||
#[case::p384("secp384r1", Some(KeyExchangeGroup::Secp384r1))]
|
||||
fn ecdh_curve_selects_the_single_key_exchange_group(
|
||||
#[case] curve: &str,
|
||||
#[case] expected: Option<KeyExchangeGroup>,
|
||||
) {
|
||||
let settings = HttpSettings {
|
||||
ssl_ecdh_curve: Some(curve.into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let resolution = Resolution::from(&settings);
|
||||
assert_eq!(resolution.config.key_exchange_group, expected);
|
||||
assert_eq!(resolution.unsupported, []);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported() {
|
||||
let settings = HttpSettings {
|
||||
ssl_ecdh_curve: Some("secp521r1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let resolution = Resolution::from(&settings);
|
||||
assert_eq!(resolution.config.key_exchange_group, None);
|
||||
assert_eq!(
|
||||
resolution.unsupported,
|
||||
[Unsupported::EcdhCurve("secp521r1".into())]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn legacy_security_level_keeps_every_suite_and_is_reported_unsupported() {
|
||||
let settings = HttpSettings {
|
||||
ssl_security_level: Some("DEFAULT@SECLEVEL=1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let resolution = Resolution::from(&settings);
|
||||
assert_eq!(resolution.config.tls12_cipher_suites, None);
|
||||
assert_eq!(
|
||||
resolution.unsupported,
|
||||
[Unsupported::SecurityLevel("@SECLEVEL=1".into())]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn named_suites_restrict_tls12_and_unsupported_entries_are_reported() {
|
||||
let settings = HttpSettings {
|
||||
ssl_security_level: Some(
|
||||
"ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES128-GCM-SHA256:!aNULL:AES256-SHA@SECLEVEL=2"
|
||||
.into(),
|
||||
),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let resolution = Resolution::from(&settings);
|
||||
assert_eq!(
|
||||
resolution.config.tls12_cipher_suites,
|
||||
Some(vec![
|
||||
Tls12CipherSuite::EcdheEcdsaAes128Gcm,
|
||||
Tls12CipherSuite::EcdheRsaAes256Gcm
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
resolution.unsupported,
|
||||
[
|
||||
Unsupported::CipherToken("!aNULL".into()),
|
||||
Unsupported::CipherToken("AES256-SHA".into())
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_settings_carry_over_unchanged() {
|
||||
let keepalive = TcpKeepalive {
|
||||
idle: Duration::from_secs(60),
|
||||
interval: Duration::from_secs(30),
|
||||
retries: 5,
|
||||
};
|
||||
let settings = HttpSettings {
|
||||
ssl_certificate: Some("/client.pem".into()),
|
||||
force_ipv4: true,
|
||||
http2: true,
|
||||
user_agent: Some("litellm/1.0".into()),
|
||||
trust_proxy_env: true,
|
||||
connect_timeout: Duration::from_secs(7),
|
||||
tcp_keepalive: Some(keepalive),
|
||||
pool_idle_timeout: Duration::from_secs(45),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let config = Resolution::from(&settings).config;
|
||||
assert_eq!(
|
||||
config,
|
||||
HttpClientConfig {
|
||||
verify: Verify::BuiltInRoots,
|
||||
client_certificate: Some("/client.pem".into()),
|
||||
key_exchange_group: None,
|
||||
tls12_cipher_suites: None,
|
||||
force_ipv4: true,
|
||||
http2: true,
|
||||
user_agent: Some("litellm/1.0".into()),
|
||||
trust_proxy_env: true,
|
||||
connect_timeout: Duration::from_secs(7),
|
||||
tcp_keepalive: Some(keepalive),
|
||||
pool_idle_timeout: Duration::from_secs(45),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_ca_bundle_is_a_read_error() {
|
||||
let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem");
|
||||
let config = HttpClientConfig {
|
||||
verify: Verify::CaBundle(path.clone()),
|
||||
..Resolution::from(&HttpSettings::default()).config
|
||||
};
|
||||
assert!(matches!(
|
||||
reqwest::ClientBuilder::try_from(&config),
|
||||
Err(Error::Read { path: reported, .. }) if reported == path
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_pem_ca_bundle_is_an_invalid_pem_error() {
|
||||
let path =
|
||||
std::env::temp_dir().join(format!("litellm-http-not-pem-{}.pem", std::process::id()));
|
||||
std::fs::write(&path, b"not a certificate").unwrap();
|
||||
let config = HttpClientConfig {
|
||||
verify: Verify::CaBundle(path.clone()),
|
||||
..Resolution::from(&HttpSettings::default()).config
|
||||
};
|
||||
let result = reqwest::ClientBuilder::try_from(&config).map(drop);
|
||||
std::fs::remove_file(&path).unwrap();
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidPem { path: reported, .. }) if reported == path
|
||||
));
|
||||
}
|
||||
}
|
||||
17
litellm-rust/crates/http/src/error.rs
Normal file
17
litellm-rust/crates/http/src/error.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum Error {
|
||||
#[error("could not read {}: {message}", path.display())]
|
||||
Read { path: PathBuf, message: String },
|
||||
#[error("{} is not a PEM file: {message}", path.display())]
|
||||
InvalidPem { path: PathBuf, message: String },
|
||||
#[error("could not build the HTTP client: {0}")]
|
||||
Client(String),
|
||||
}
|
||||
|
||||
impl From<reqwest::Error> for Error {
|
||||
fn from(error: reqwest::Error) -> Self {
|
||||
Self::Client(error.without_url().to_string())
|
||||
}
|
||||
}
|
||||
13
litellm-rust/crates/http/src/lib.rs
Normal file
13
litellm-rust/crates/http/src/lib.rs
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
mod config;
|
||||
mod error;
|
||||
mod pool;
|
||||
mod proxy;
|
||||
mod settings;
|
||||
mod tls;
|
||||
|
||||
pub use config::{HttpClientConfig, Resolution, Verify};
|
||||
pub use error::Error;
|
||||
pub use pool::{ClientVariant, HttpClientPool};
|
||||
pub use proxy::EnvironmentProxies;
|
||||
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};
|
||||
pub use tls::{KeyExchangeGroup, Tls12CipherSuite, Unsupported};
|
||||
325
litellm-rust/crates/http/src/pool.rs
Normal file
325
litellm-rust/crates/http/src/pool.rs
Normal file
|
|
@ -0,0 +1,325 @@
|
|||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, Mutex, MutexGuard, PoisonError},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use reqwest::dns::Resolve;
|
||||
|
||||
use crate::{config::HttpClientConfig, error::Error};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum ClientVariant {
|
||||
Provider,
|
||||
NoRedirect,
|
||||
Media,
|
||||
UnpinnedMedia,
|
||||
}
|
||||
|
||||
const CLIENT_TTL: Duration = Duration::from_secs(3600);
|
||||
|
||||
struct PooledClient {
|
||||
client: reqwest::Client,
|
||||
built_at: Instant,
|
||||
}
|
||||
|
||||
type Clients = HashMap<(HttpClientConfig, ClientVariant), PooledClient>;
|
||||
|
||||
pub struct HttpClientPool {
|
||||
media_resolver: Arc<dyn Resolve>,
|
||||
ttl: Duration,
|
||||
clients: Mutex<Clients>,
|
||||
}
|
||||
|
||||
impl HttpClientPool {
|
||||
pub fn new(media_resolver: Arc<dyn Resolve>) -> Self {
|
||||
Self::with_ttl(media_resolver, CLIENT_TTL)
|
||||
}
|
||||
|
||||
pub fn with_ttl(media_resolver: Arc<dyn Resolve>, ttl: Duration) -> Self {
|
||||
Self {
|
||||
media_resolver,
|
||||
ttl,
|
||||
clients: Mutex::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn client(
|
||||
&self,
|
||||
config: &HttpClientConfig,
|
||||
variant: ClientVariant,
|
||||
) -> Result<reqwest::Client, Error> {
|
||||
let effective = match variant {
|
||||
ClientVariant::Media => HttpClientConfig {
|
||||
client_certificate: None,
|
||||
trust_proxy_env: false,
|
||||
..config.clone()
|
||||
},
|
||||
ClientVariant::UnpinnedMedia => HttpClientConfig {
|
||||
client_certificate: None,
|
||||
..config.clone()
|
||||
},
|
||||
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
|
||||
};
|
||||
let key = (effective, variant);
|
||||
if let Some(pooled) = self.lock().get(&key)
|
||||
&& pooled.built_at.elapsed() < self.ttl
|
||||
{
|
||||
return Ok(pooled.client.clone());
|
||||
}
|
||||
let client = self
|
||||
.apply(variant, reqwest::ClientBuilder::try_from(&key.0)?)
|
||||
.build()?;
|
||||
self.lock().insert(
|
||||
key,
|
||||
PooledClient {
|
||||
client: client.clone(),
|
||||
built_at: Instant::now(),
|
||||
},
|
||||
);
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
fn lock(&self) -> MutexGuard<'_, Clients> {
|
||||
self.clients.lock().unwrap_or_else(PoisonError::into_inner)
|
||||
}
|
||||
|
||||
fn apply(
|
||||
&self,
|
||||
variant: ClientVariant,
|
||||
builder: reqwest::ClientBuilder,
|
||||
) -> reqwest::ClientBuilder {
|
||||
match variant {
|
||||
ClientVariant::Provider => builder,
|
||||
ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => {
|
||||
builder.redirect(reqwest::redirect::Policy::none())
|
||||
}
|
||||
ClientVariant::Media => builder
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.dns_resolver2(Arc::clone(&self.media_resolver)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{
|
||||
net::SocketAddr,
|
||||
sync::atomic::{AtomicUsize, Ordering},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use reqwest::dns::{Addrs, Name, Resolving};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpListener,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
use crate::{HttpSettings, Resolution, Verify};
|
||||
|
||||
struct FixedResolver(SocketAddr);
|
||||
|
||||
impl Resolve for FixedResolver {
|
||||
fn resolve(&self, _: Name) -> Resolving {
|
||||
let addrs: Addrs = Box::new(std::iter::once(self.0));
|
||||
Box::pin(std::future::ready(Ok(addrs)))
|
||||
}
|
||||
}
|
||||
|
||||
fn pool() -> HttpClientPool {
|
||||
HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())))
|
||||
}
|
||||
|
||||
fn config(user_agent: &str) -> HttpClientConfig {
|
||||
HttpClientConfig {
|
||||
user_agent: Some(user_agent.into()),
|
||||
..Resolution::from(&HttpSettings::default()).config
|
||||
}
|
||||
}
|
||||
|
||||
async fn serve(
|
||||
status_line: &'static str,
|
||||
) -> (SocketAddr, Arc<AtomicUsize>, Arc<Mutex<Vec<String>>>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let connections = Arc::new(AtomicUsize::new(0));
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let (accepted, seen) = (Arc::clone(&connections), Arc::clone(&requests));
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
accepted.fetch_add(1, Ordering::SeqCst);
|
||||
let seen = Arc::clone(&seen);
|
||||
tokio::spawn(async move {
|
||||
let mut buffer = vec![0u8; 4096];
|
||||
while let Ok(read) = socket.read(&mut buffer).await {
|
||||
if read == 0 {
|
||||
return;
|
||||
}
|
||||
seen.lock()
|
||||
.unwrap()
|
||||
.push(String::from_utf8_lossy(&buffer[..read]).into_owned());
|
||||
let response = format!(
|
||||
"{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\n\r\n"
|
||||
);
|
||||
if socket.write_all(response.as_bytes()).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
(address, connections, requests)
|
||||
}
|
||||
|
||||
async fn get(
|
||||
pool: &HttpClientPool,
|
||||
config: &HttpClientConfig,
|
||||
variant: ClientVariant,
|
||||
url: &str,
|
||||
) -> reqwest::Response {
|
||||
pool.client(config, variant)
|
||||
.unwrap()
|
||||
.get(url)
|
||||
.timeout(Duration::from_secs(5))
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn clients_are_shared_per_config_and_variant() {
|
||||
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
|
||||
let url = format!("http://{address}");
|
||||
let pool = pool();
|
||||
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
|
||||
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
|
||||
assert_eq!(connections.load(Ordering::SeqCst), 1);
|
||||
get(&pool, &config("a"), ClientVariant::NoRedirect, &url).await;
|
||||
assert_eq!(connections.load(Ordering::SeqCst), 2);
|
||||
get(&pool, &config("b"), ClientVariant::Provider, &url).await;
|
||||
assert_eq!(connections.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expired_clients_are_rebuilt() {
|
||||
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
|
||||
let url = format!("http://{address}");
|
||||
let pool = HttpClientPool::with_ttl(
|
||||
Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())),
|
||||
Duration::ZERO,
|
||||
);
|
||||
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
|
||||
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
|
||||
assert_eq!(connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn media_clients_are_shared_across_proxy_settings_they_never_use() {
|
||||
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
|
||||
let pool = HttpClientPool::new(Arc::new(FixedResolver(address)));
|
||||
let url = format!("http://media.invalid:{}/doc", address.port());
|
||||
for trust_proxy_env in [true, false] {
|
||||
let config = HttpClientConfig {
|
||||
trust_proxy_env,
|
||||
..config("a")
|
||||
};
|
||||
get(&pool, &config, ClientVariant::Media, &url).await;
|
||||
}
|
||||
assert_eq!(connections.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn media_variant_never_loads_the_client_certificate() {
|
||||
let pool = pool();
|
||||
let with_identity = HttpClientConfig {
|
||||
client_certificate: Some(std::env::temp_dir().join("litellm-http-absent-client.pem")),
|
||||
..config("a")
|
||||
};
|
||||
assert!(
|
||||
pool.client(&with_identity, ClientVariant::Provider)
|
||||
.is_err()
|
||||
);
|
||||
assert!(pool.client(&with_identity, ClientVariant::Media).is_ok());
|
||||
assert!(
|
||||
pool.client(&with_identity, ClientVariant::UnpinnedMedia)
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_failures_are_not_cached() {
|
||||
let pool = pool();
|
||||
let missing = HttpClientConfig {
|
||||
verify: Verify::CaBundle(std::env::temp_dir().join("litellm-http-absent.pem")),
|
||||
..config("a")
|
||||
};
|
||||
assert!(pool.client(&missing, ClientVariant::Provider).is_err());
|
||||
assert!(pool.client(&missing, ClientVariant::Provider).is_err());
|
||||
assert!(pool.client(&config("a"), ClientVariant::Provider).is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_client_sends_the_configured_user_agent_over_http1() {
|
||||
let (address, _, requests) = serve("HTTP/1.1 204 No Content").await;
|
||||
let response = get(
|
||||
&pool(),
|
||||
&config("litellm-test/9"),
|
||||
ClientVariant::Provider,
|
||||
&format!("http://{address}"),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 204);
|
||||
assert_eq!(response.version(), reqwest::Version::HTTP_11);
|
||||
let request = requests.lock().unwrap()[0].clone();
|
||||
assert!(request.contains("user-agent: litellm-test/9"), "{request}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() {
|
||||
let (address, _, _) = serve("HTTP/1.1 302 Found").await;
|
||||
let response = get(
|
||||
&pool(),
|
||||
&config("a"),
|
||||
ClientVariant::NoRedirect,
|
||||
&format!("http://{address}"),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 302);
|
||||
assert_eq!(response.headers()["location"], "/elsewhere");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unpinned_media_variant_uses_the_system_resolver_and_returns_redirects() {
|
||||
let (address, _, _) = serve("HTTP/1.1 302 Found").await;
|
||||
let pool = HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())));
|
||||
let response = get(
|
||||
&pool,
|
||||
&config("a"),
|
||||
ClientVariant::UnpinnedMedia,
|
||||
&format!("http://localhost:{}/doc", address.port()),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 302);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn media_variant_resolves_through_the_injected_resolver() {
|
||||
let (address, _, requests) = serve("HTTP/1.1 204 No Content").await;
|
||||
let pool = HttpClientPool::new(Arc::new(FixedResolver(address)));
|
||||
let url = format!("http://media.invalid:{}/doc", address.port());
|
||||
let response = get(&pool, &config("a"), ClientVariant::Media, &url).await;
|
||||
assert_eq!(response.status(), 204);
|
||||
assert!(requests.lock().unwrap()[0].contains("host: media.invalid"));
|
||||
assert!(
|
||||
pool.client(&config("a"), ClientVariant::Provider)
|
||||
.unwrap()
|
||||
.get(&url)
|
||||
.timeout(Duration::from_secs(5))
|
||||
.send()
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
}
|
||||
15
litellm-rust/crates/http/src/proxy.rs
Normal file
15
litellm-rust/crates/http/src/proxy.rs
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
use hyper_util::client::proxy::matcher::Matcher;
|
||||
|
||||
pub struct EnvironmentProxies(Matcher);
|
||||
|
||||
impl EnvironmentProxies {
|
||||
pub fn from_environment() -> Self {
|
||||
Self(Matcher::from_system())
|
||||
}
|
||||
|
||||
pub fn apply_to(&self, url: &reqwest::Url) -> bool {
|
||||
url.as_str()
|
||||
.parse::<http::Uri>()
|
||||
.is_ok_and(|uri| self.0.intercept(&uri).is_some())
|
||||
}
|
||||
}
|
||||
416
litellm-rust/crates/http/src/settings.rs
Normal file
416
litellm-rust/crates/http/src/settings.rs
Normal file
|
|
@ -0,0 +1,416 @@
|
|||
use std::{
|
||||
path::{Path, PathBuf},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum SslVerify {
|
||||
Enabled,
|
||||
Disabled,
|
||||
CaBundle(PathBuf),
|
||||
}
|
||||
|
||||
impl SslVerify {
|
||||
pub fn parse(value: &str) -> Self {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"true" => Self::Enabled,
|
||||
"false" => Self::Disabled,
|
||||
_ => Self::CaBundle(PathBuf::from(value)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct TcpKeepalive {
|
||||
pub idle: Duration,
|
||||
pub interval: Duration,
|
||||
pub retries: u32,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct HttpSettingsLayer {
|
||||
pub ssl_verify: Option<SslVerify>,
|
||||
pub ssl_cert_file: Option<PathBuf>,
|
||||
pub ssl_certificate: Option<PathBuf>,
|
||||
pub ssl_security_level: Option<String>,
|
||||
pub ssl_ecdh_curve: Option<String>,
|
||||
pub force_ipv4: Option<bool>,
|
||||
pub http2: Option<bool>,
|
||||
pub aiohttp_trust_env: Option<bool>,
|
||||
pub disable_aiohttp_trust_env: Option<bool>,
|
||||
pub disable_aiohttp_transport: Option<bool>,
|
||||
pub user_agent: Option<String>,
|
||||
pub tcp_keepalive: Option<TcpKeepalive>,
|
||||
pub pool_idle_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl HttpSettingsLayer {
|
||||
pub fn from_environment(env: &(dyn Fn(&str) -> Option<String> + Sync)) -> Self {
|
||||
let enabled = |name: &str| {
|
||||
env(name)
|
||||
.is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
|
||||
.then_some(true)
|
||||
};
|
||||
let number = |name: &str| env(name).and_then(|value| value.trim().parse::<u32>().ok());
|
||||
let seconds = |name: &str, default: u32| {
|
||||
Duration::from_secs(u64::from(number(name).unwrap_or(default)))
|
||||
};
|
||||
Self {
|
||||
ssl_verify: env("SSL_VERIFY").map(|value| SslVerify::parse(&value)),
|
||||
ssl_cert_file: env("SSL_CERT_FILE").map(PathBuf::from),
|
||||
ssl_certificate: env("SSL_CERTIFICATE").map(PathBuf::from),
|
||||
ssl_security_level: env("SSL_SECURITY_LEVEL"),
|
||||
ssl_ecdh_curve: env("SSL_ECDH_CURVE"),
|
||||
force_ipv4: None,
|
||||
http2: enabled("LITELLM_HTTP2"),
|
||||
aiohttp_trust_env: enabled("AIOHTTP_TRUST_ENV"),
|
||||
disable_aiohttp_trust_env: enabled("DISABLE_AIOHTTP_TRUST_ENV"),
|
||||
disable_aiohttp_transport: enabled("DISABLE_AIOHTTP_TRANSPORT"),
|
||||
user_agent: env("LITELLM_USER_AGENT"),
|
||||
tcp_keepalive: enabled("AIOHTTP_SO_KEEPALIVE").map(|_| TcpKeepalive {
|
||||
idle: seconds("AIOHTTP_TCP_KEEPIDLE", 60),
|
||||
interval: seconds("AIOHTTP_TCP_KEEPINTVL", 30),
|
||||
retries: number("AIOHTTP_TCP_KEEPCNT").unwrap_or(5),
|
||||
}),
|
||||
pool_idle_timeout: number("AIOHTTP_KEEPALIVE_TIMEOUT")
|
||||
.map(|timeout| Duration::from_secs(u64::from(timeout))),
|
||||
}
|
||||
}
|
||||
|
||||
fn or(self, lower: Self) -> Self {
|
||||
Self {
|
||||
ssl_verify: self.ssl_verify.or(lower.ssl_verify),
|
||||
ssl_cert_file: self.ssl_cert_file.or(lower.ssl_cert_file),
|
||||
ssl_certificate: self.ssl_certificate.or(lower.ssl_certificate),
|
||||
ssl_security_level: self.ssl_security_level.or(lower.ssl_security_level),
|
||||
ssl_ecdh_curve: self.ssl_ecdh_curve.or(lower.ssl_ecdh_curve),
|
||||
force_ipv4: self.force_ipv4.or(lower.force_ipv4),
|
||||
http2: self.http2.or(lower.http2),
|
||||
aiohttp_trust_env: self.aiohttp_trust_env.or(lower.aiohttp_trust_env),
|
||||
disable_aiohttp_trust_env: self
|
||||
.disable_aiohttp_trust_env
|
||||
.or(lower.disable_aiohttp_trust_env),
|
||||
disable_aiohttp_transport: self
|
||||
.disable_aiohttp_transport
|
||||
.or(lower.disable_aiohttp_transport),
|
||||
user_agent: self.user_agent.or(lower.user_agent),
|
||||
tcp_keepalive: self.tcp_keepalive.or(lower.tcp_keepalive),
|
||||
pool_idle_timeout: self.pool_idle_timeout.or(lower.pool_idle_timeout),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct HttpSettings {
|
||||
pub ssl_verify: Option<SslVerify>,
|
||||
pub ssl_cert_file: Option<PathBuf>,
|
||||
pub ssl_certificate: Option<PathBuf>,
|
||||
pub ssl_security_level: Option<String>,
|
||||
pub ssl_ecdh_curve: Option<String>,
|
||||
pub force_ipv4: bool,
|
||||
pub http2: bool,
|
||||
pub user_agent: Option<String>,
|
||||
pub trust_proxy_env: bool,
|
||||
pub connect_timeout: Duration,
|
||||
pub tcp_keepalive: Option<TcpKeepalive>,
|
||||
pub pool_idle_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for HttpSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
ssl_verify: None,
|
||||
ssl_cert_file: None,
|
||||
ssl_certificate: None,
|
||||
ssl_security_level: None,
|
||||
ssl_ecdh_curve: None,
|
||||
force_ipv4: false,
|
||||
http2: false,
|
||||
user_agent: None,
|
||||
trust_proxy_env: true,
|
||||
connect_timeout: Duration::from_secs(10),
|
||||
tcp_keepalive: None,
|
||||
pool_idle_timeout: Duration::from_secs(120),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpSettings {
|
||||
pub fn from_layers(
|
||||
highest_precedence_first: impl IntoIterator<Item = HttpSettingsLayer>,
|
||||
) -> Self {
|
||||
let merged = highest_precedence_first
|
||||
.into_iter()
|
||||
.reduce(HttpSettingsLayer::or)
|
||||
.unwrap_or_default();
|
||||
let defaults = Self::default();
|
||||
let http2 = merged.http2.unwrap_or(defaults.http2);
|
||||
Self {
|
||||
ssl_verify: merged.ssl_verify,
|
||||
ssl_cert_file: merged.ssl_cert_file,
|
||||
ssl_certificate: merged
|
||||
.ssl_certificate
|
||||
.filter(|path| !path.as_os_str().is_empty()),
|
||||
ssl_security_level: merged.ssl_security_level.filter(|level| !level.is_empty()),
|
||||
ssl_ecdh_curve: merged.ssl_ecdh_curve.filter(|curve| !curve.is_empty()),
|
||||
force_ipv4: merged.force_ipv4.unwrap_or(defaults.force_ipv4),
|
||||
http2,
|
||||
user_agent: merged.user_agent,
|
||||
trust_proxy_env: !merged.disable_aiohttp_trust_env.unwrap_or(false)
|
||||
|| merged.aiohttp_trust_env.unwrap_or(false)
|
||||
|| merged.disable_aiohttp_transport.unwrap_or(false)
|
||||
|| http2,
|
||||
tcp_keepalive: merged.tcp_keepalive,
|
||||
pool_idle_timeout: merged
|
||||
.pool_idle_timeout
|
||||
.unwrap_or(defaults.pool_idle_timeout),
|
||||
..defaults
|
||||
}
|
||||
}
|
||||
|
||||
pub fn without_missing_files(self, exists: &dyn Fn(&Path) -> bool) -> Self {
|
||||
Self {
|
||||
ssl_verify: match self.ssl_verify {
|
||||
Some(SslVerify::CaBundle(path)) if !exists(&path) => Some(SslVerify::Enabled),
|
||||
other => other,
|
||||
},
|
||||
ssl_cert_file: self.ssl_cert_file.filter(|path| exists(path)),
|
||||
..self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn no_env(_: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn env_of(
|
||||
values: &'static [(&'static str, &'static str)],
|
||||
) -> impl Fn(&str) -> Option<String> + Sync {
|
||||
move |name| {
|
||||
values
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| value.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("true", SslVerify::Enabled)]
|
||||
#[case(" True ", SslVerify::Enabled)]
|
||||
#[case("FALSE", SslVerify::Disabled)]
|
||||
#[case("/etc/ssl/bundle.pem", SslVerify::CaBundle("/etc/ssl/bundle.pem".into()))]
|
||||
fn ssl_verify_parses_bools_and_treats_anything_else_as_a_bundle_path(
|
||||
#[case] value: &str,
|
||||
#[case] expected: SslVerify,
|
||||
) {
|
||||
assert_eq!(SslVerify::parse(value), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn higher_layers_override_lower_ones() {
|
||||
let configured = HttpSettingsLayer {
|
||||
ssl_verify: Some(SslVerify::Enabled),
|
||||
ssl_certificate: Some("/configured/client.pem".into()),
|
||||
ssl_security_level: Some("configured".into()),
|
||||
user_agent: Some("configured/1".into()),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let environment = HttpSettingsLayer::from_environment(&env_of(&[
|
||||
("SSL_VERIFY", "false"),
|
||||
("SSL_CERT_FILE", "/env/roots.pem"),
|
||||
("SSL_CERTIFICATE", "/env/client.pem"),
|
||||
("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1"),
|
||||
("SSL_ECDH_CURVE", "X25519"),
|
||||
("LITELLM_USER_AGENT", "env/2"),
|
||||
]));
|
||||
let settings = HttpSettings::from_layers([environment, configured]);
|
||||
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
|
||||
assert_eq!(settings.ssl_cert_file, Some("/env/roots.pem".into()));
|
||||
assert_eq!(settings.ssl_certificate, Some("/env/client.pem".into()));
|
||||
assert_eq!(
|
||||
settings.ssl_security_level.as_deref(),
|
||||
Some("DEFAULT@SECLEVEL=1")
|
||||
);
|
||||
assert_eq!(settings.ssl_ecdh_curve.as_deref(), Some("X25519"));
|
||||
assert_eq!(settings.user_agent.as_deref(), Some("env/2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_explicit_false_in_a_higher_layer_beats_a_lower_true() {
|
||||
let higher = HttpSettingsLayer {
|
||||
http2: Some(false),
|
||||
force_ipv4: Some(false),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let lower = HttpSettingsLayer {
|
||||
http2: Some(true),
|
||||
force_ipv4: Some(true),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let settings = HttpSettings::from_layers([higher, lower]);
|
||||
assert!(!settings.http2);
|
||||
assert!(!settings.force_ipv4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_empty_environment_is_an_empty_layer_so_lower_layers_and_defaults_apply() {
|
||||
assert_eq!(
|
||||
HttpSettingsLayer::from_environment(&no_env),
|
||||
HttpSettingsLayer::default()
|
||||
);
|
||||
let configured = HttpSettingsLayer {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
|
||||
http2: Some(true),
|
||||
user_agent: Some("configured/1".into()),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
assert_eq!(
|
||||
HttpSettings::from_layers([HttpSettingsLayer::default(), configured]),
|
||||
HttpSettings {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
|
||||
http2: true,
|
||||
user_agent: Some("configured/1".into()),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
);
|
||||
assert_eq!(HttpSettings::from_layers([]), HttpSettings::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_environment_values_clear_the_setting_like_python_truthiness() {
|
||||
let configured = HttpSettingsLayer {
|
||||
ssl_certificate: Some("/configured/client.pem".into()),
|
||||
ssl_security_level: Some("configured".into()),
|
||||
ssl_ecdh_curve: Some("X25519".into()),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let environment = HttpSettingsLayer::from_environment(&env_of(&[
|
||||
("SSL_CERTIFICATE", ""),
|
||||
("SSL_SECURITY_LEVEL", ""),
|
||||
("SSL_ECDH_CURVE", ""),
|
||||
]));
|
||||
let settings = HttpSettings::from_layers([environment, configured]);
|
||||
assert_eq!(settings.ssl_certificate, None);
|
||||
assert_eq!(settings.ssl_security_level, None);
|
||||
assert_eq!(settings.ssl_ecdh_curve, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn socket_keepalive_follows_the_aiohttp_variables_with_python_defaults() {
|
||||
let tuned = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(&[
|
||||
("AIOHTTP_SO_KEEPALIVE", "True"),
|
||||
("AIOHTTP_TCP_KEEPIDLE", "45"),
|
||||
("AIOHTTP_KEEPALIVE_TIMEOUT", "30"),
|
||||
]))]);
|
||||
assert_eq!(
|
||||
tuned.tcp_keepalive,
|
||||
Some(TcpKeepalive {
|
||||
idle: Duration::from_secs(45),
|
||||
interval: Duration::from_secs(30),
|
||||
retries: 5,
|
||||
})
|
||||
);
|
||||
assert_eq!(tuned.pool_idle_timeout, Duration::from_secs(30));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn socket_keepalive_stays_off_unless_enabled() {
|
||||
let settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(
|
||||
&[("AIOHTTP_TCP_KEEPIDLE", "45")],
|
||||
))]);
|
||||
assert_eq!(settings.tcp_keepalive, None);
|
||||
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
|
||||
}
|
||||
|
||||
fn proxy_flags(
|
||||
aiohttp_trust_env: bool,
|
||||
disable_aiohttp_trust_env: bool,
|
||||
disable_aiohttp_transport: bool,
|
||||
http2: bool,
|
||||
) -> HttpSettingsLayer {
|
||||
HttpSettingsLayer {
|
||||
aiohttp_trust_env: Some(aiohttp_trust_env),
|
||||
disable_aiohttp_trust_env: Some(disable_aiohttp_trust_env),
|
||||
disable_aiohttp_transport: Some(disable_aiohttp_transport),
|
||||
http2: Some(http2),
|
||||
..HttpSettingsLayer::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::aiohttp_default(proxy_flags(false, false, false, false), true)]
|
||||
#[case::aiohttp_opted_out(proxy_flags(false, true, false, false), false)]
|
||||
#[case::session_trust_env_beats_opt_out(proxy_flags(true, true, false, false), true)]
|
||||
#[case::http2_uses_httpx(proxy_flags(false, true, false, true), true)]
|
||||
#[case::aiohttp_disabled(proxy_flags(false, true, true, false), true)]
|
||||
fn environment_proxies_apply_unless_the_aiohttp_transport_opts_out(
|
||||
#[case] layer: HttpSettingsLayer,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
assert_eq!(HttpSettings::from_layers([layer]).trust_proxy_env, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_proxy_opt_out_in_one_source_still_yields_to_trust_env_from_another() {
|
||||
let environment =
|
||||
HttpSettingsLayer::from_environment(&env_of(&[("DISABLE_AIOHTTP_TRUST_ENV", "true")]));
|
||||
let configured = HttpSettingsLayer {
|
||||
aiohttp_trust_env: Some(true),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
assert!(!HttpSettings::from_layers([environment.clone()]).trust_proxy_env);
|
||||
assert!(HttpSettings::from_layers([environment, configured]).trust_proxy_env);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_files_fall_back_to_default_verification() {
|
||||
let settings = HttpSettings {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/absent/roots.pem".into())),
|
||||
ssl_cert_file: Some("/absent/env.pem".into()),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
.without_missing_files(&|_| false);
|
||||
assert_eq!(settings.ssl_verify, Some(SslVerify::Enabled));
|
||||
assert_eq!(settings.ssl_cert_file, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn existing_files_are_kept() {
|
||||
let settings = HttpSettings {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/present/roots.pem".into())),
|
||||
ssl_cert_file: Some("/present/env.pem".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
assert_eq!(settings.clone().without_missing_files(&|_| true), settings);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("true", Some(true))]
|
||||
#[case("True", Some(true))]
|
||||
#[case("false", None)]
|
||||
#[case("1", None)]
|
||||
fn boolean_switches_only_turn_on_for_true(
|
||||
#[case] value: &'static str,
|
||||
#[case] expected: Option<bool>,
|
||||
) {
|
||||
let env = move |name: &str| match name {
|
||||
"LITELLM_HTTP2"
|
||||
| "AIOHTTP_TRUST_ENV"
|
||||
| "DISABLE_AIOHTTP_TRANSPORT"
|
||||
| "DISABLE_AIOHTTP_TRUST_ENV" => Some(value.to_string()),
|
||||
_ => None,
|
||||
};
|
||||
let layer = HttpSettingsLayer::from_environment(&env);
|
||||
assert_eq!(layer.http2, expected);
|
||||
assert_eq!(layer.aiohttp_trust_env, expected);
|
||||
assert_eq!(layer.disable_aiohttp_transport, expected);
|
||||
assert_eq!(layer.disable_aiohttp_trust_env, expected);
|
||||
}
|
||||
}
|
||||
411
litellm-rust/crates/http/src/tls.rs
Normal file
411
litellm-rust/crates/http/src/tls.rs
Normal 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
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -20,6 +20,8 @@ bytes.workspace = true
|
|||
litellm-auth.workspace = true
|
||||
litellm-callbacks-legacy.workspace = true
|
||||
litellm-core.workspace = true
|
||||
litellm-auth-gcp.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-types.workspace = true
|
||||
litellm-host-python.workspace = true
|
||||
|
|
|
|||
18
litellm-rust/crates/python-bridge/python_settings.json
Normal file
18
litellm-rust/crates/python-bridge/python_settings.json
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
{
|
||||
"http_settings": [
|
||||
"ssl_verify",
|
||||
"ssl_certificate",
|
||||
"ssl_security_level",
|
||||
"ssl_ecdh_curve",
|
||||
"force_ipv4",
|
||||
"http2",
|
||||
"aiohttp_trust_env",
|
||||
"disable_aiohttp_trust_env",
|
||||
"disable_aiohttp_transport",
|
||||
"user_agent"
|
||||
],
|
||||
"url_policy": [
|
||||
"user_url_validation",
|
||||
"user_url_allowed_hosts"
|
||||
]
|
||||
}
|
||||
354
litellm-rust/crates/python-bridge/src/http.rs
Normal file
354
litellm-rust/crates/python-bridge/src/http.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,9 @@
|
|||
mod credentials;
|
||||
mod diagnostics;
|
||||
mod errors;
|
||||
mod http;
|
||||
mod marshal;
|
||||
mod python_settings;
|
||||
mod routes;
|
||||
mod token_counter;
|
||||
|
||||
|
|
|
|||
65
litellm-rust/crates/python-bridge/src/python_settings.rs
Normal file
65
litellm-rust/crates/python-bridge/src/python_settings.rs
Normal 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);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -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 },
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 = "",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -52,6 +52,15 @@ class BaseAudioTranscriptionConfig(BaseConfig, ABC):
|
|||
"""
|
||||
return False
|
||||
|
||||
@property
|
||||
def has_native_transcription_endpoint(self) -> bool:
|
||||
"""
|
||||
Opt-in for OpenAI-compatible providers whose transcription lives on a
|
||||
non-OpenAI route: when True the request skips the OpenAI SDK transport
|
||||
and goes through this config via the shared http handler.
|
||||
"""
|
||||
return False
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
|
|
|
|||
275
litellm/llms/base_llm/realtime/transcription_protocol.py
Normal file
275
litellm/llms/base_llm/realtime/transcription_protocol.py
Normal 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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -2223,6 +2223,7 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
# Prepare headers
|
||||
kwargs = kwargs or {}
|
||||
kwargs_for_agentic: Final = self._agentic_hook_kwargs(kwargs=kwargs, api_key=api_key, api_base=api_base)
|
||||
provider_specific_header: Final = cast(
|
||||
litellm.types.utils.ProviderSpecificHeader | Sequence[litellm.types.utils.ProviderSpecificHeader] | None,
|
||||
kwargs.get("provider_specific_header", None),
|
||||
|
|
@ -2410,7 +2411,7 @@ class BaseLLMHTTPHandler:
|
|||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
|
||||
kwargs=kwargs_for_agentic,
|
||||
hold_back=bool(held_back_tool_names),
|
||||
server_fulfilled_tool_names=held_back_tool_names,
|
||||
)
|
||||
|
|
@ -2433,8 +2434,7 @@ class BaseLLMHTTPHandler:
|
|||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_key=api_key,
|
||||
kwargs=kwargs,
|
||||
kwargs=kwargs_for_agentic,
|
||||
)
|
||||
|
||||
async def _finalize_anthropic_messages_response(
|
||||
|
|
@ -2447,14 +2447,8 @@ class BaseLLMHTTPHandler:
|
|||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
custom_llm_provider: str,
|
||||
api_key: str | None,
|
||||
kwargs: dict,
|
||||
kwargs: dict[str, object],
|
||||
) -> AnthropicMessagesResponse | AsyncIterator:
|
||||
# Inject api_key into kwargs so follow-up calls in agentic hooks can
|
||||
# authenticate. api_key is a named param here (not in kwargs), so
|
||||
# _prepare_followup_kwargs would miss it otherwise.
|
||||
kwargs_for_agentic: Final = {**kwargs, "api_key": api_key} if api_key else kwargs
|
||||
# Call agentic completion hooks (non-streaming path only)
|
||||
final_response: Final = await self._call_agentic_completion_hooks(
|
||||
response=initial_response,
|
||||
model=model,
|
||||
|
|
@ -2464,7 +2458,7 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj=logging_obj,
|
||||
stream=False,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs_for_agentic,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
return self._maybe_wrap_in_fake_stream(
|
||||
|
|
@ -5312,6 +5306,15 @@ class BaseLLMHTTPHandler:
|
|||
fingerprints: Final = list(kwargs.get("_agentic_loop_fingerprints", []) or [])
|
||||
return depth, max_loops, fingerprints
|
||||
|
||||
@staticmethod
|
||||
def _agentic_hook_kwargs(
|
||||
kwargs: Mapping[str, object], api_key: str | None, api_base: str | None
|
||||
) -> dict[str, object]:
|
||||
"""``api_key`` and ``api_base`` are named parameters of ``anthropic_messages`` rather than kwargs, so the
|
||||
follow-up call an agentic hook makes only reaches the same deployment if they are re-added here."""
|
||||
deployment_params: Final = {"api_key": api_key, "api_base": api_base}
|
||||
return {**kwargs, **{key: value for key, value in deployment_params.items() if value}}
|
||||
|
||||
@staticmethod
|
||||
def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool:
|
||||
"""
|
||||
|
|
@ -5906,7 +5909,15 @@ class BaseLLMHTTPHandler:
|
|||
callback.__class__.__name__,
|
||||
plan.stop_reason,
|
||||
)
|
||||
return self._maybe_wrap_in_fake_stream(response, logging_obj, api_surface)
|
||||
return self._maybe_wrap_in_fake_stream(
|
||||
await callback.async_post_agentic_loop_response_hook(
|
||||
response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls),
|
||||
plan=plan,
|
||||
kwargs=kwargs_with_provider,
|
||||
),
|
||||
logging_obj,
|
||||
api_surface,
|
||||
)
|
||||
if not plan.run_agentic_loop:
|
||||
continue
|
||||
|
||||
|
|
@ -6279,7 +6290,12 @@ class BaseLLMHTTPHandler:
|
|||
ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
||||
ssl_context.check_hostname = False
|
||||
ssl_context.verify_mode = ssl.CERT_NONE
|
||||
backend_ws: Final = await self._open_realtime_backend_ws(websockets, url, headers, ssl_context)
|
||||
provider_backend: Final = await provider_config.open_backend(url, headers)
|
||||
backend_ws: Final = (
|
||||
provider_backend
|
||||
if provider_backend is not None
|
||||
else await self._open_realtime_backend_ws(websockets, url, headers, ssl_context)
|
||||
)
|
||||
async with backend_ws:
|
||||
_request_data: Final[dict[str, object]] = {}
|
||||
if litellm_metadata:
|
||||
|
|
@ -6591,7 +6607,7 @@ class BaseLLMHTTPHandler:
|
|||
first_message: str | None = None,
|
||||
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
) -> Exception | None:
|
||||
"""
|
||||
Handles Responses API WebSocket mode.
|
||||
|
||||
|
|
@ -6625,7 +6641,7 @@ class BaseLLMHTTPHandler:
|
|||
**kwargs,
|
||||
)
|
||||
await handler.run()
|
||||
return
|
||||
return None
|
||||
|
||||
import websockets
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
|
|
@ -6744,9 +6760,10 @@ class BaseLLMHTTPHandler:
|
|||
output_guardrail_callbacks=_ws_output_guardrail_callbacks,
|
||||
quota_callbacks=_ws_quota_callbacks,
|
||||
authorized_model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
request_defaults=request_defaults,
|
||||
)
|
||||
await streaming.bidirectional_forward()
|
||||
return await streaming.bidirectional_forward()
|
||||
|
||||
except websockets.exceptions.InvalidStatusCode as e:
|
||||
verbose_logger.exception("Error connecting to responses WS backend: %s", e)
|
||||
|
|
@ -6760,6 +6777,7 @@ class BaseLLMHTTPHandler:
|
|||
pass
|
||||
else:
|
||||
raise Exception(f"Unexpected error while closing WebSocket: {close_error}")
|
||||
return None
|
||||
|
||||
def image_edit_handler(
|
||||
self,
|
||||
|
|
|
|||
42
litellm/llms/fireworks_ai/cache_pricing.py
Normal file
42
litellm/llms/fireworks_ai/cache_pricing.py
Normal 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
|
||||
),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
418
litellm/llms/vertex_ai/audio_transcription/realtime_backend.py
Normal file
418
litellm/llms/vertex_ai/audio_transcription/realtime_backend.py
Normal 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))
|
||||
|
|
@ -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]
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -12,10 +12,16 @@ Auth: OAuth2 Bearer token (not an API key).
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Final
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import (
|
||||
VertexChirpRealtimeConfig,
|
||||
is_vertex_speech_to_text_model,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
|
||||
class VertexAIRealtimeConfig(GeminiRealtimeConfig):
|
||||
|
|
@ -232,3 +238,20 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
|
|||
return []
|
||||
|
||||
return super().transform_realtime_request(message, model, session_configuration_request)
|
||||
|
||||
|
||||
def vertex_realtime_config(
|
||||
model: str,
|
||||
*,
|
||||
access_token: str,
|
||||
resolve_access_token: Callable[[], Awaitable[str]],
|
||||
project: str,
|
||||
location: str | None,
|
||||
) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig:
|
||||
if is_vertex_speech_to_text_model(model):
|
||||
return VertexChirpRealtimeConfig(resolve_access_token=resolve_access_token, project=project, location=location)
|
||||
return VertexAIRealtimeConfig(
|
||||
access_token=access_token,
|
||||
project=project,
|
||||
location=VertexBase.get_vertex_region(vertex_region=location, model=model),
|
||||
)
|
||||
|
|
|
|||
3
litellm/llms/xai/audio_transcription/__init__.py
Normal file
3
litellm/llms/xai/audio_transcription/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .transformation import XAIAudioTranscriptionConfig
|
||||
|
||||
__all__ = ["XAIAudioTranscriptionConfig"]
|
||||
207
litellm/llms/xai/audio_transcription/transformation.py
Normal file
207
litellm/llms/xai/audio_transcription/transformation.py
Normal 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}"}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -51,7 +51,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
|
|||
from fastapi import HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
|
|
@ -187,6 +187,52 @@ class MintProxyCredential(Protocol):
|
|||
) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ...
|
||||
|
||||
|
||||
TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
|
||||
|
||||
|
||||
def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]:
|
||||
"""The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed
|
||||
only where the JWT auth that proves a subject token is on, backed by a database, and
|
||||
licensed, so a client never selects a grant the gateway would then refuse."""
|
||||
if token_exchange_available:
|
||||
return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE)
|
||||
return ("authorization_code", "refresh_token")
|
||||
|
||||
|
||||
"""RFC 8693: a native client that already holds a token from the customer's identity
|
||||
provider trades it for the proxy-API credential without a browser round trip."""
|
||||
|
||||
_IssuedTokenType = Literal["urn:ietf:params:oauth:token-type:access_token"]
|
||||
ACCESS_TOKEN_TOKEN_TYPE: Final[_IssuedTokenType] = "urn:ietf:params:oauth:token-type:access_token"
|
||||
SUBJECT_TOKEN_TYPES: Final = frozenset(
|
||||
{
|
||||
"urn:ietf:params:oauth:token-type:jwt",
|
||||
"urn:ietf:params:oauth:token-type:id_token",
|
||||
ACCESS_TOKEN_TOKEN_TYPE,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class SubjectIdentity(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
user_id: str = Field(min_length=1)
|
||||
team_id: str | None = None
|
||||
|
||||
|
||||
class SubjectTokenRefusal(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
error: Literal["unsupported_grant_type", "invalid_request", "temporarily_unavailable"]
|
||||
description: str = Field(min_length=1)
|
||||
|
||||
|
||||
class ExchangeSubjectToken(Protocol):
|
||||
"""Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the
|
||||
IdP token the way the proxy's own JWT auth does and names the litellm user and team it
|
||||
stands for, or says why this gateway will not take it."""
|
||||
|
||||
def __call__(self, subject_token: str, request: Request, /) -> Awaitable[SubjectIdentity | SubjectTokenRefusal]: ...
|
||||
|
||||
|
||||
class ConsentTeam(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
team_id: str = Field(min_length=1)
|
||||
|
|
@ -213,6 +259,12 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr
|
|||
return "unresolvable"
|
||||
|
||||
|
||||
async def _refuse_subject_token(subject_token: str, request: Request) -> SubjectTokenRefusal:
|
||||
return SubjectTokenRefusal(
|
||||
error="unsupported_grant_type", description="this gateway is not configured to exchange IdP tokens"
|
||||
)
|
||||
|
||||
|
||||
async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState:
|
||||
return "unavailable"
|
||||
|
||||
|
|
@ -318,7 +370,9 @@ def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None:
|
|||
return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY)
|
||||
|
||||
|
||||
async def register_aggregate_client(request: Request, request_body: Mapping[str, object]) -> Response:
|
||||
async def register_aggregate_client(
|
||||
request: Request, request_body: Mapping[str, object], token_exchange_available: bool
|
||||
) -> Response:
|
||||
"""RFC 7591 dynamic registration against the gateway itself, statelessly.
|
||||
|
||||
Only ``redirect_uris`` is authoritative; every client is registered as a public
|
||||
|
|
@ -382,7 +436,7 @@ async def register_aggregate_client(request: Request, request_body: Mapping[str,
|
|||
"client_id_issued_at": int(now.timestamp()),
|
||||
"redirect_uris": list(raw_uris),
|
||||
"token_endpoint_auth_method": "none",
|
||||
"grant_types": ["authorization_code", "refresh_token"],
|
||||
"grant_types": list(supported_grant_types(token_exchange_available)),
|
||||
"response_types": ["code"],
|
||||
},
|
||||
)
|
||||
|
|
@ -580,7 +634,7 @@ class NativeClientAuthContract(TypedDict):
|
|||
revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]]
|
||||
|
||||
|
||||
def native_client_auth_contract(request: Request) -> NativeClientAuthContract:
|
||||
def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract:
|
||||
"""The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a
|
||||
native client (in any language) needs to run the sign-in without reading LiteLLM
|
||||
source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter
|
||||
|
|
@ -595,7 +649,7 @@ def native_client_auth_contract(request: Request) -> NativeClientAuthContract:
|
|||
"revocation_endpoint": f"{base_url}/revoke",
|
||||
"resource": base_url,
|
||||
"response_types_supported": ("code",),
|
||||
"grant_types_supported": ("authorization_code", "refresh_token"),
|
||||
"grant_types_supported": supported_grant_types(token_exchange_available),
|
||||
"code_challenge_methods_supported": ("S256",),
|
||||
"token_endpoint_auth_methods_supported": ("none",),
|
||||
"revocation_endpoint_auth_methods_supported": ("none",),
|
||||
|
|
@ -1033,20 +1087,26 @@ class _ProxyCredentialTokenResponse(TypedDict):
|
|||
refresh_token: ReadOnly[str]
|
||||
user_id: ReadOnly[str]
|
||||
team_id: ReadOnly[str | None]
|
||||
issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]]
|
||||
|
||||
|
||||
def _proxy_credential_response(
|
||||
minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime
|
||||
minted: MintedProxyCredential,
|
||||
principal: SessionPrincipal,
|
||||
keys: SessionSigningKeys,
|
||||
now: datetime,
|
||||
issued_token_type: _IssuedTokenType | None = None,
|
||||
) -> Response:
|
||||
"""The proxy-API token response: the access token is the very credential ``lite
|
||||
login`` stores (accepted on every proxy route with user and team attribution), and
|
||||
the refresh token is a gateway-sealed rotating token bound to the team the credential
|
||||
was minted for, so a renewal keeps the team the user consented to."""
|
||||
was minted for, so a renewal keeps the team the user consented to. A token exchange
|
||||
also states ``issued_token_type``, which RFC 8693 section 2.2.1 requires."""
|
||||
bound_principal: Final = principal.model_copy(update=MappingProxyType({"team_id": minted.team_id}))
|
||||
refresh: Final = mint_session_refresh_token(bound_principal, keys, now)
|
||||
if not isinstance(refresh, MintedSessionToken):
|
||||
return _oauth_error(500, "server_error", "failed to mint the session credential")
|
||||
body: Final[_ProxyCredentialTokenResponse] = {
|
||||
credential: Final[_ProxyCredentialTokenResponse] = {
|
||||
"access_token": minted.key,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": minted.expires_in,
|
||||
|
|
@ -1054,7 +1114,10 @@ def _proxy_credential_response(
|
|||
"user_id": minted.user_id,
|
||||
"team_id": minted.team_id,
|
||||
}
|
||||
return JSONResponse(status_code=200, content=body, headers=TOKEN_NO_CACHE_HEADERS)
|
||||
if issued_token_type is None:
|
||||
return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS)
|
||||
exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type}
|
||||
return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS)
|
||||
|
||||
|
||||
def _reload_failure_response(failure: ReloadUserFailure) -> Response:
|
||||
|
|
@ -1073,6 +1136,16 @@ def _reload_failure_response(failure: ReloadUserFailure) -> Response:
|
|||
assert_never(failure)
|
||||
|
||||
|
||||
def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response:
|
||||
match refusal.error:
|
||||
case "temporarily_unavailable":
|
||||
return _oauth_error(503, refusal.error, refusal.description)
|
||||
case "unsupported_grant_type" | "invalid_request":
|
||||
return _oauth_error(400, refusal.error, refusal.description)
|
||||
case _:
|
||||
assert_never(refusal.error)
|
||||
|
||||
|
||||
def _mint_failure_response(failure: ProxyCredentialMintFailure) -> Response:
|
||||
match failure:
|
||||
case "not_a_member":
|
||||
|
|
@ -1116,11 +1189,16 @@ async def aggregate_token(
|
|||
cache: DualCache,
|
||||
resource: str | None = None,
|
||||
mint_proxy_credential: MintProxyCredential = _refuse_proxy_credential,
|
||||
subject_token: str | None = None,
|
||||
subject_token_type: str | None = None,
|
||||
requested_token_type: str | None = None,
|
||||
exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token,
|
||||
) -> Response:
|
||||
"""The aggregate token verb: authorization_code and refresh_token grants for the
|
||||
identity-only session pair, or for the proxy-API credential when the grant was issued
|
||||
with that audience. Every path re-validates the litellm user live before minting, so a
|
||||
deactivated user cannot obtain or renew a session."""
|
||||
with that audience, and the RFC 8693 token exchange that turns an IdP token straight
|
||||
into the proxy-API credential. Every path re-validates the litellm user live before
|
||||
minting, so a deactivated user cannot obtain or renew a session."""
|
||||
if master_key is None:
|
||||
verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured")
|
||||
return _oauth_error(500, "server_error", "the gateway has no master key configured")
|
||||
|
|
@ -1159,7 +1237,20 @@ async def aggregate_token(
|
|||
now=now,
|
||||
issue=issue,
|
||||
)
|
||||
return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token")
|
||||
if grant_type == TOKEN_EXCHANGE_GRANT_TYPE:
|
||||
return await _token_exchange_grant(
|
||||
subject_token=subject_token,
|
||||
subject_token_type=subject_token_type,
|
||||
requested_token_type=requested_token_type,
|
||||
client_id=client_id,
|
||||
exchange_subject_token=exchange_subject_token,
|
||||
issue=issue,
|
||||
)
|
||||
return _oauth_error(
|
||||
400,
|
||||
"unsupported_grant_type",
|
||||
f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}",
|
||||
)
|
||||
|
||||
|
||||
class _GrantIssuer:
|
||||
|
|
@ -1211,10 +1302,9 @@ class _GrantIssuer:
|
|||
async def _issue_proxy_credential(
|
||||
self, principal: SessionPrincipal, claim_key: str, claim_ttl_seconds: int, replayed: str
|
||||
) -> Response:
|
||||
if self._resource is not None and not is_proxy_api_resource(self._request, self._resource):
|
||||
return _oauth_error(
|
||||
400, "invalid_target", "resource does not match the proxy API this grant was issued for"
|
||||
)
|
||||
target_refusal: Final = self._proxy_api_target_refusal()
|
||||
if target_refusal is not None:
|
||||
return target_refusal
|
||||
minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id)
|
||||
if not isinstance(minted, MintedProxyCredential):
|
||||
return _mint_failure_response(minted)
|
||||
|
|
@ -1223,6 +1313,33 @@ class _GrantIssuer:
|
|||
return refusal
|
||||
return _proxy_credential_response(minted, principal, self._keys, self._now)
|
||||
|
||||
async def exchange(
|
||||
self, subject_token: str, client_id: str, exchange_subject_token: ExchangeSubjectToken
|
||||
) -> Response:
|
||||
"""The RFC 8693 tail: prove the IdP token, then mint. No single-use marker, because
|
||||
the subject token stays a valid proof for as long as the IdP says it is and every
|
||||
exchange mints a fresh credential and refresh token of its own."""
|
||||
target_refusal: Final = self._proxy_api_target_refusal()
|
||||
if target_refusal is not None:
|
||||
return target_refusal
|
||||
identity: Final = await exchange_subject_token(subject_token, self._request)
|
||||
if isinstance(identity, SubjectTokenRefusal):
|
||||
return _subject_token_refusal_response(identity)
|
||||
principal: Final = SessionPrincipal(
|
||||
user_id=identity.user_id, client_id=client_id, audience=PROXY_API_AUDIENCE, team_id=identity.team_id
|
||||
)
|
||||
minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id)
|
||||
if not isinstance(minted, MintedProxyCredential):
|
||||
return _mint_failure_response(minted)
|
||||
return _proxy_credential_response(
|
||||
minted, principal, self._keys, self._now, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE
|
||||
)
|
||||
|
||||
def _proxy_api_target_refusal(self) -> Response | None:
|
||||
if self._resource is None or is_proxy_api_resource(self._request, self._resource):
|
||||
return None
|
||||
return _oauth_error(400, "invalid_target", "resource does not match the proxy API this grant was issued for")
|
||||
|
||||
async def _claim_refusal(self, claim_key: str, claim_ttl_seconds: int, replayed: str) -> Response | None:
|
||||
return _claim_refusal(
|
||||
await self._guard.claim(claim_key, claim_ttl_seconds), replayed=_oauth_error(400, "invalid_grant", replayed)
|
||||
|
|
@ -1297,6 +1414,32 @@ async def _refresh_token_grant(
|
|||
)
|
||||
|
||||
|
||||
async def _token_exchange_grant(
|
||||
subject_token: str | None,
|
||||
subject_token_type: str | None,
|
||||
requested_token_type: str | None,
|
||||
client_id: str,
|
||||
exchange_subject_token: ExchangeSubjectToken,
|
||||
issue: _GrantIssuer,
|
||||
) -> Response:
|
||||
"""RFC 8693 token exchange for a registered native client that already holds an IdP
|
||||
token: the gateway proves the token the way its JWT auth does and answers with the
|
||||
proxy-API credential, so a fresh laptop with only an IdP login gets a gateway key
|
||||
without a browser round trip. The client must be registered because the refresh token
|
||||
in the answer is bound to it."""
|
||||
if not is_gateway_dcr_client_id(client_id) or open_gateway_dcr_client(client_id) is None:
|
||||
return _oauth_error(401, "invalid_client", "unknown or malformed client_id")
|
||||
if not subject_token or not subject_token_type:
|
||||
return _oauth_error(400, "invalid_request", "subject_token and subject_token_type are required")
|
||||
if subject_token_type not in SUBJECT_TOKEN_TYPES:
|
||||
return _oauth_error(
|
||||
400, "invalid_request", f"subject_token_type must be one of {', '.join(sorted(SUBJECT_TOKEN_TYPES))}"
|
||||
)
|
||||
if requested_token_type is not None and requested_token_type != ACCESS_TOKEN_TOKEN_TYPE:
|
||||
return _oauth_error(400, "invalid_request", f"requested_token_type must be {ACCESS_TOKEN_TOKEN_TYPE}")
|
||||
return await issue.exchange(subject_token, client_id, exchange_subject_token)
|
||||
|
||||
|
||||
async def revoke_refresh_token(token: str, client_id: str, master_key: str | None, cache: DualCache) -> Response:
|
||||
"""RFC 7009 revocation for the gateway's refresh tokens: burn the presented token's
|
||||
``jti`` so neither the holder nor a thief can rotate it again. Access tokens are
|
||||
|
|
|
|||
217
litellm/proxy/_experimental/mcp_server/idp_token_exchange.py
Normal file
217
litellm/proxy/_experimental/mcp_server/idp_token_exchange.py
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
"""The identity-provider side of the RFC 8693 token exchange on ``POST /token``: a native
|
||||
client that already holds a JWT from the customer's IdP trades it for the same proxy-API
|
||||
credential ``lite login`` stores, proven by the proxy's own JWT auth (signature, claims,
|
||||
and the user and team sync it performs), so no browser round trip is needed."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal
|
||||
from litellm.proxy._types import JWTAuthBuilderResult, ProxyException
|
||||
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
|
||||
EXCHANGE_ROUTE: Final = "/token"
|
||||
REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth"
|
||||
SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = (
|
||||
"the gateway could not verify subject_token because its identity provider or database is unavailable; retry"
|
||||
)
|
||||
SUBJECT_TOKEN_CHECK_FAULTED: Final = (
|
||||
"the gateway could not verify subject_token because its database reported a fault that is not a transient "
|
||||
"outage; retrying will not help until the gateway deployment is repaired"
|
||||
)
|
||||
GatewayOutage = Literal["retryable", "faulted"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenExchangePrerequisites:
|
||||
"""The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT
|
||||
bearer, plus the JWT-to-virtual-key mapping it consults first: a gateway that maps
|
||||
tokens authenticates a JWT as its mapped key, with that key's models and budget, or
|
||||
refuses an unmapped one, and the exchange proves the token through ``auth_builder``
|
||||
alone, so it would mint the user's own credential past that policy. Discovery and
|
||||
registration advertise the exchange grant only when every gate holds, and an exchange
|
||||
attempt is refused naming the first one that does not."""
|
||||
|
||||
jwt_auth_enabled: bool
|
||||
has_database: bool
|
||||
licensed: bool
|
||||
maps_jwts_to_virtual_keys: bool
|
||||
|
||||
@property
|
||||
def available(self) -> bool:
|
||||
return self.jwt_auth_enabled and self.has_database and self.licensed and not self.maps_jwts_to_virtual_keys
|
||||
|
||||
def refusal(self) -> SubjectTokenRefusal | None:
|
||||
if not self.jwt_auth_enabled:
|
||||
return SubjectTokenRefusal(
|
||||
error="unsupported_grant_type",
|
||||
description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens",
|
||||
)
|
||||
if not self.has_database:
|
||||
return SubjectTokenRefusal(
|
||||
error="unsupported_grant_type",
|
||||
description="this gateway has no database, so it cannot exchange IdP tokens",
|
||||
)
|
||||
if not self.licensed:
|
||||
return SubjectTokenRefusal(
|
||||
error="unsupported_grant_type",
|
||||
description="JWT auth is an enterprise only feature; no license is set",
|
||||
)
|
||||
if self.maps_jwts_to_virtual_keys:
|
||||
return SubjectTokenRefusal(
|
||||
error="unsupported_grant_type",
|
||||
description="this gateway maps IdP tokens to virtual keys, which the exchange does not serve",
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def read_token_exchange_prerequisites() -> TokenExchangePrerequisites:
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
|
||||
general_settings,
|
||||
jwt_handler,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
)
|
||||
|
||||
return TokenExchangePrerequisites(
|
||||
jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True,
|
||||
has_database=prisma_client is not None,
|
||||
licensed=premium_user is True,
|
||||
maps_jwts_to_virtual_keys=_maps_jwts_to_virtual_keys(jwt_handler),
|
||||
)
|
||||
|
||||
|
||||
def _maps_jwts_to_virtual_keys(jwt_handler: JWTHandler) -> bool:
|
||||
if not hasattr(jwt_handler, "litellm_jwtauth"):
|
||||
return False
|
||||
return jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured()
|
||||
|
||||
|
||||
def token_exchange_available() -> bool:
|
||||
return read_token_exchange_prerequisites().available
|
||||
|
||||
|
||||
class AuthorizeSubjectToken(Protocol):
|
||||
"""Injected JWT authorization ``(subject_token, request_headers)``: the proxy's
|
||||
``JWTAuthManager.auth_builder`` in production, which raises when the token is not
|
||||
acceptable and otherwise names the user and team it resolved."""
|
||||
|
||||
def __call__(
|
||||
self, subject_token: str, request_headers: Mapping[str, str], /
|
||||
) -> Awaitable[JWTAuthBuilderResult]: ...
|
||||
|
||||
|
||||
async def exchange_idp_subject_token(subject_token: str, request: Request) -> SubjectIdentity | SubjectTokenRefusal:
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
|
||||
general_settings,
|
||||
jwt_handler,
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
async def authorize(token: str, request_headers: Mapping[str, str]) -> JWTAuthBuilderResult:
|
||||
return await JWTAuthManager.auth_builder(
|
||||
api_key=token,
|
||||
jwt_handler=jwt_handler,
|
||||
request_data={},
|
||||
general_settings=general_settings,
|
||||
route=EXCHANGE_ROUTE,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
request_headers=request_headers,
|
||||
request_method="POST",
|
||||
)
|
||||
|
||||
return await identity_from_subject_token(
|
||||
subject_token,
|
||||
request_headers=request.headers,
|
||||
prerequisites=read_token_exchange_prerequisites(),
|
||||
is_jwt=jwt_handler.is_jwt,
|
||||
authorize=authorize,
|
||||
)
|
||||
|
||||
|
||||
async def identity_from_subject_token(
|
||||
subject_token: str,
|
||||
request_headers: Mapping[str, str],
|
||||
prerequisites: TokenExchangePrerequisites,
|
||||
is_jwt: Callable[[str], bool],
|
||||
authorize: AuthorizeSubjectToken,
|
||||
) -> SubjectIdentity | SubjectTokenRefusal:
|
||||
"""Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the
|
||||
proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which
|
||||
RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token, and a
|
||||
check the gateway could not complete (the IdP's JWKS unreachable with no cached copy,
|
||||
the auth database down) as ``temporarily_unavailable``, so the client retries instead
|
||||
of treating a valid token as bad, worded by whether retrying can help. The reason stays
|
||||
in the proxy log: this endpoint is public and JWT auth's own wording can name the JWKS
|
||||
URL it fetched or quote the IdP's response."""
|
||||
unmet: Final = prerequisites.refusal()
|
||||
if unmet is not None:
|
||||
return unmet
|
||||
if not is_jwt(subject_token):
|
||||
return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT")
|
||||
try:
|
||||
result: Final = await authorize(subject_token, request_headers)
|
||||
except HTTPException as denied:
|
||||
return _refusal_for(denied, denied.detail)
|
||||
except ProxyException as denied:
|
||||
return _refusal_for(denied, denied.message)
|
||||
except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures
|
||||
return _refusal_for(denied, denied)
|
||||
user_id: Final = result["user_id"]
|
||||
if user_id is None:
|
||||
return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows")
|
||||
return SubjectIdentity(user_id=user_id, team_id=result["team_id"])
|
||||
|
||||
|
||||
def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal:
|
||||
outage: Final = _gateway_could_not_verify(denied)
|
||||
if outage is None:
|
||||
verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason)
|
||||
return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
|
||||
verbose_proxy_logger.error("token exchange could not verify a subject_token, %s: %s", outage, reason)
|
||||
return SubjectTokenRefusal(error="temporarily_unavailable", description=_check_unavailable_description(outage))
|
||||
|
||||
|
||||
def _check_unavailable_description(outage: GatewayOutage) -> str:
|
||||
match outage:
|
||||
case "retryable":
|
||||
return SUBJECT_TOKEN_CHECK_UNAVAILABLE
|
||||
case "faulted":
|
||||
return SUBJECT_TOKEN_CHECK_FAULTED
|
||||
case _:
|
||||
assert_never(outage)
|
||||
|
||||
|
||||
def _gateway_could_not_verify(denied: Exception) -> GatewayOutage | None:
|
||||
"""A database fault anywhere in the chain (``get_user_object`` wraps prisma failures in a
|
||||
bare ``ValueError``) or a 5xx from JWT auth (the IdP's JWKS unreachable with no cached
|
||||
copy) is the gateway failing, not the token. A fault retrying cannot clear (a missing or
|
||||
version-skewed query engine) is named as such, the way the mint path words it, so the
|
||||
client is not told to wait on a deployment that needs repair."""
|
||||
fault: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(denied)
|
||||
if fault is not None:
|
||||
return "faulted" if PrismaDBExceptionHandler.is_permanent_database_fault(fault) else "retryable"
|
||||
return "retryable" if _is_server_error(denied) else None
|
||||
|
||||
|
||||
def _is_server_error(denied: Exception) -> bool:
|
||||
match denied:
|
||||
case HTTPException(status_code=status_code):
|
||||
return status_code >= 500
|
||||
case ProxyException(code=code):
|
||||
return code.isdigit() and int(code) >= 500
|
||||
case _:
|
||||
return False
|
||||
|
|
@ -16,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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
397
litellm/proxy/anthropic_endpoints/gateway_endpoints.py
Normal file
397
litellm/proxy/anthropic_endpoints/gateway_endpoints.py
Normal 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()
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)}."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
57
litellm/rust_bridge/settings.py
Normal file
57
litellm/rust_bridge/settings.py
Normal 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
Loading…
Add table
Reference in a new issue